This repo aims for Multi-turn Reinforcement Learning for Vision Language Models.
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
## TRICO Algorithm
VAGEN introduces TRICO, a specialized algorithm for visual agents that extends previous approaches with several key improvements:
Building on recent advances in LLM reinforcement learning, TRICO explores several key innovations for visual agents:
***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
These examples were cherry-picked from validation results. Notably, **Trico 6** (LM + GM + Turn-wise Reward + Bi-Level GAE) demonstrates enhanced capability for agent exploration and potential to solve harder problems, while **Trico 3** (LM + GM) shows more stability in solving simpler one-step and two-step problems.
-**Bi-Level-GAE**: Whether to use multi-turn GAE (first estimates turn-level advantage, then estimates advantage in each turn)
-`algorithm.adv_estimator=bi_level_gae`
-**Turn-Wise-GAE**: Whether to use turn-aware GAE (each turn will have only one same advantage estimation)
-`algorithm.adv_estimator=turn_wise_gae`
-**Masked-GAE**: Whether to use masked GAE (skips observation tokens from environment when estimating advantages)
-`algorithm.adv_estimator=masked_gae`
### Configuration Settings
-**multi-turn-reward**: Whether to use multi-turn reward (gives step reward for last token of each turn, instead of summing all rewards for last token of whole trajectory)
-`rollout_manager.use_multi_turn_reward=True`
-**with_loss_mask**: Whether to use loss mask to only calculate the loss of tokens output by the models
-`rollout_manager.use_loss_mask=True`
## Conclusion
VAGEN, powered by the TRICO algorithm, is well-suited for visual agent learning because:
1.**Visual agents require multi-turn reasoning**: The agent observes, thinks, acts, and receives updated observations
2.**Not all tokens are equally important**: In visual reasoning, the action tokens are more critical than observation tokens for VLMs to learn
3.**Bi-level time dependencies**: Bi-Level discount factors help model both immediate and long-term consequences
## Future Work
1. Expanding to more environments, real-world applications like GUI agents and embodied agents
