@@ -38,40 +38,14 @@ Traditional RL frameworks for LLM agents treat all tokens in a trajectory equall
VAGEN addresses these challenges by focusing optimization on the most critical decision-making tokens and creating a more nuanced reward structure across interaction turns.
## Experimental Results
Our experiments on visual Sokoban using a Qwen-VL 3B model show:
- TRICO significantly outperforms RICO in visual agentic tasks
- Both selective token masking and cross-turn credit assignment contribute to performance gains
- AICO (Action-centric Interaction Chain Optimization), which uses only selective token masking, outperforms TRICO on simple tasks
- TRICO demonstrates superior exploration capabilities on more complex problems
bash vagen/examples/frozen_lake_aico/run.sh # AICO without service
bash vagen/examples/frozen_lake_aico_service/run.sh # AICO with service
bash vagen/examples/frozen_lake_trico/run.sh # TRICO without service
# SVG Generation
bash vagen/examples/svg_aico/run.sh # AICO without service
bash vagen/examples/svg_trico/run.sh # TRICO without service
```
Each run takes ~4 hours to reach 150 steps on 4 H100s. You can decrease testing frequency to speed up training. Training might be unstable due to loss spikes; we recommend restoring from the latest checkpoint when encountering such cases. We will resolve this issue in future work (see roadmap).
## How to Add New Environment
## Service and Environment Architecture
We introduce a new service-based architecture that addresses the challenges of managing multiple embodied environments for VLM agent training. This design enables efficient parallel processing across distributed systems, providing better scalability, standardization, and the ability to seamlessly integrate both rule-based rewards and reward models within the same framework
VAGEN supports creating custom environments for agent training, which could simply inherit from `BaseEnv` and `BaseEnvConfig` classes to implement your environment.
### Directory Structure
```
vagen/
├── env/
│ ├── base_service.py # Abstract base class defining the service interface
│ ├── client.py # Client for interacting with environment server
│ ├── server.py # Server implementation for hosting environments
│ └── REGISTERED_ENV # Registry mapping environment names to services
├── utils/
| ├── serial.py # Handle observation serilization from service to client
```
For detailed instructions, see our [Creating Environments](docs/create-env.md) guide. You may also want to check our [Creating Service](docs/create-service.md) for scaling your environments.
### Component Hierarchy
```
BaseService (ABC)
├── **Batch Methods**
├── create_environments_batch()
├── reset_batch()
├── step_batch()
├── compute_reward_batch()
├── get_system_prompts_batch()
└── close_batch()
BatchEnvClient
├── **HTTP Communication**
├── **Batch Methods**
└── **Convenience Methods**
BatchEnvServer
├── **Service Management**
├── **Request Routing**
├── **Batch Method Implementation**
└── **Server Management**
```
## How to Add New Model
VAGEN supports integration with various language models. To add a new model:
### New Service Development
To develop a new environment service, you only need to:
1. Define model interface in the mllm_agent architecture
2. Implement model-specific adapters and handlers
3. Configure the model in your training scripts
- Create a service class that inherits from `BaseService`
- Register the service with `REGISTERED_ENV`
For detailed implementation examples, refer to our code architecture documentation based on the [VERL architecture](https://verl.readthedocs.io/en/latest/index.html).
The FrozenLake service example demonstrates how to:
- Handle batch operations
## Experimental Results
> To reproduce our experiment, please refer to document: [Reproduce Experiments](docs/reproduce-exp.md)
The SVG service example demonstrates how to:
- Integrate Reward Models (like DINO) directly within the service
- Combine rule-based rewards with model-based rewards
## Algorithm Settings
Our experiments on visual Sokoban using a Qwen-VL 3B model show:
- TRICO significantly outperforms RICO in visual agentic tasks
- Both selective token masking and cross-turn credit assignment contribute to performance gains
- AICO (Action-centric Interaction Chain Optimization), which uses only selective token masking, outperforms TRICO on simple tasks
- TRICO demonstrates superior exploration capabilities on more complex problems
-**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`
-**with_loss_mask**: Whether to use gae mask to only calculate the gae of tokens output by the models
-`rollout_manager.use_gae_mask=True`
## Cases
We present several cases selected from validation steps during training models with AICO and TRICO, as shown below. You can view all the cases in our [Experiment Log](https://api.wandb.ai/links/ragen-V/nlb40e7l).
