VAGEN is a multi-turn Reinforcement Learning (RL) framework designed to optimize Vision-Language Models (VLMs) for visual reasoning tasks. The framework introduces the **Turn-aware Reason-Interaction Chain Optimization (TRICO)** algorithm, which enhances VLM performance on complex visual reasoning problems.
### Key Innovations
-**Selective Training**: Uses masks to control which tokens contribute to the learning process
-**Selective Reward Assignment**: Assigns turn-wise or trajectory rewards to specific tokens
-**Selective Advantage Calculation**: Implements different advantage calculations for within-turn vs. cross-turn dependencies
VAGEN is a multi-turn reinforcement learning framework designed specifically for training Vision-Language Model (VLM) Agents. Building upon existing approaches for LLM agents like RAGEN, Search-R1, and Agent-R1, VAGEN introduces enhancements that better handle the unique challenges of visual agents.
## TRICO Overview
## Key Innovations
Building on recent advances in LLM reinforcement learning, we propose TRICO, a specialized algorithm for visual agents with serval explorations, as shown in figures below.
VAGEN introduces the **Turn-aware Reason-Interaction Chain Optimization (TRICO)** algorithm which extends the traditional RICO approach with two key innovations:
1.**Selective Token Masking** - Focuses optimization on action-critical tokens through:
- Loss masking (`M^loss`): Identifies tokens to update during policy optimization
- Advantage masking (`M^adv`): Determines tokens to include in advantage calculations
***Selective Training**: Uses masks to control which tokens contribute to the learning process and advantage estimation process
***LM**: Loss Mask, state tokens (image tokens and text description tokens) are masked during the critic and actor update
***GM**: GAE Mask, state tokens are masked during generalized advantage estimation (GAE)
***Selective Reward Assignment**: Assign turn-wise reward or trajectory reward to tokens
***Turn-wise reward**: Reward is assigned to each turn's <eos> token
***Trajectory reward**: Reward is summed up and assigned to the last token of the trajectory
***Selective Advantage Calculation**: Different advantage calculation for within-turn vs. cross-turn dependencies
***Bi-level GAE**: We use two discounting factors:
* $γ_{turn}$: Estimates advantage for each turn through cross-turn temporal difference learning
* $γ_{token}$: Estimates advantage for each token through within-turn temporal difference learning
***Turn-Wise GAE**: We use $γ_{turn}$ to compute a single advantage through cross-turn temporal difference learning and assign the same advantage for each token in the turn
2.**Cross-turn Credit Assignment** - Enables more effective credit attribution through:
- Bi-level advantage estimation with separate discount factors for cross-turn (`γ_turn`) and within-turn (`γ_token`) calculations
- Turn-level rewards applied at each interaction boundary
## Why VAGEN Works Better for VLM Agents
### Algorithm Features
Traditional RL frameworks for LLM agents treat all tokens in a trajectory equally. This approach is suboptimal for VLM agents due to:
-**Distribution Shift**: Most VLMs aren't pretrained to generate image tokens
-**State Redundancy**: Visual tasks contain excessive low-level information in long-context inputs
## Performance
VAGEN addresses these challenges by focusing optimization on the most critical decision-making tokens and creating a more nuanced reward structure across interaction turns.
We evaluated VAGEN on the visual puzzle-solving Sokoban task, demonstrating significant improvements over previous approaches.
| **LM output** | No special structure | `<think>...</think><ans>...</ans>` | `<think>...</think><ans>...</ans><eoa>` |
| **Discounting** | Single discount rate | Single discount rate | Bi-level discounting |
| **Optimization** | All tokens equally | All tokens equally | Selective token optimization |
The visualizations below show how the agent reasons through sequential steps to solve Sokoban puzzles, cherry picked from validation steps when training TRICO: LM + GM and TRICO: LM + GM + Turn-wise Reward + Bi-Level GAE.
