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 RL to build explicit internal world models by structuring reasoning into StateEstimation (grounding current visual observations) and TransitionModeling (predicting next states). A 3B-parameter model trained with VAGEN-Full achieves 0.82 overall across five diverse visual agent tasks, outperforming GPT-5 (0.75), Gemini 2.5 Pro (0.67), and Claude 4.5 (0.62).

Problem

Existing RL frameworks for VLM agents treat multi-turn visual tasks like text-based MDPs, ignoring that visual observations are partial and noisy—making the problem a POMDP. Without explicit world modeling, VLMs fail to maintain and update internal beliefs across turns, leading to poor performance on long-horizon visual agentic tasks. Standard RL methods (Vanilla PPO, GRPO, Turn-PPO) also suffer from sparse reward propagation over long trajectories.

Method

The POMDP is formulated with five reasoning strategies enforced via format reward during RL training: NoThink, FreeThink, StateEstimation (generates <observation>ŝt</observation>), TransitionModeling (generates <prediction>ŝt+1</prediction>), and WorldModeling (both). The full VAGEN-Full system adds two mechanisms on top of WorldModeling:

  1. WorldModeling Reward: A dense, turn-by-turn reward using LLM-as-a-Judge to score the agent's state descriptions and next-state predictions against ground-truth environment states (positions, distances, coordinates), replacing the noisy CLIP-based signal.

  2. Bi-Level GAE: Two-stage credit assignment—first computing turn-level advantages via GAE across turns (discount γ_turn), then propagating each turn's advantage backward to individual tokens within that turn (discount γ_token), initialized at the final token of each action. This replaces single-pass token-level GAE that dilutes sparse end-of-trajectory rewards.

Training backbone is Qwen2.5-VL-3B using PPO via the verl library. Representation format (natural language vs. structured vs. symbolic) is selected per-task.

Key Contributions

  • Formal POMDP framing of multi-turn VLM agent training with explicit StateEstimation and TransitionModeling as structured reasoning tokens
  • Empirical comparison of five reasoning strategies showing WorldModeling dominates across diverse tasks
  • Task-dependent analysis of internal belief representations: natural language beats symbolic/structured for general tasks; structured coordinates are essential for high-precision manipulation
  • WorldModeling Reward: dense LLM-as-a-Judge turn-level reward derived from ground-truth environment state annotations
  • Bi-Level GAE: hierarchical credit assignment separating inter-turn and intra-turn advantage propagation
  • VAGEN framework: decoupled environment/training infrastructure supporting five visually distinct environments (Sokoban, FrozenLake, 3D Navigation, PrimitiveSkill robot manipulation, SVG Reconstruction)

Results

  • VAGEN-Full (3B) overall score: 0.82 vs. untrained baseline 0.21 (~3× improvement)
  • Proprietary model comparison on same benchmark: GPT-5 0.75, o3 0.73, Gemini 2.5 Pro 0.64, Claude 4.5 Sonnet 0.54
  • WorldModeling reasoning strategy: 0.76 vs. FreeThink 0.67 vs. NoThink 0.28 (VAGEN-Base)
  • VAGEN-Full vs. VAGEN-Base (WorldModeling): largest gap on PrimitiveSkill (0.97 vs. 0.91 average subtask), indicating better generalization to new scenes
  • RL baselines on same backbone: Vanilla-PPO 0.26, GRPO w/ Mask 0.54, Turn-PPO w/ Mask 0.55 vs. VAGEN-Full 0.82
  • Representation: Natural language > Structured > Symbolic for FrozenLake/Sokoban; Structured slightly > Natural language for PrimitiveSkill manipulation

Limitations

  • Evaluated on a single model family (Qwen2.5-VL-3B); generalizability across architectures unverified
  • Only five environments tested; coverage of diverse real-world visual agent scenarios is limited
  • WorldModeling Reward requires ground-truth state annotations from the environment, which may not be available in unstructured or real-world settings
  • SVG Reconstruction lacks world dynamics, so WorldModeling Reward is inapplicable there; only Bi-Level GAE is used
  • CLIP-based reward was tried and discarded due to insensitivity to fine-grained spatial detail, suggesting reward quality is environment-dependent
  • No supervised finetuning (SFT) warm-up stage; authors note future work will add multi-turn visual understanding SFT

Relevance to Agentic AI / LLM Agents

This paper directly addresses the core challenge of POMDP-structured multi-turn visual agent training—a gap that has limited VLM agents relative to text-only LLM agents. The structured world-modeling reasoning approach (StateEstimation + TransitionModeling) provides a concrete, RL-trainable formalization of belief maintenance, advancing beyond free-form chain-of-thought. The Bi-Level GAE method is applicable beyond VLMs to any multi-turn agent RL setting with dense intermediate signals, making it a broadly relevant credit assignment contribution. The finding that a 3B open-source model can surpass frontier proprietary models (GPT-5, Gemini 2.5 Pro) on visual agentic benchmarks via structured RL training is a significant capability result for the field.