Merge pull request #16 from RAGEN-AI/dev

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

<img width="800" alt="image" src="./public/1.png" />

<img width="800" alt="image" src="./public/2.png" />

<img width="800" alt="image" src="./public/3.png" />

## Comparison of Algorithms

| **Feature** | **PPO** | **RICO** | **TRICO (Ours)** |

| --- | --- | --- | --- |

| **Sequence Structure** | Single response | Multiple turn interaction | Multiple turn interaction |

| **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 |

## Training Configuration

We used the following settings in our experiments:

- **Model**: Qwen 2.5 VL-instruction 3B

- **Environment**: Visual Sokoban (puzzle-solving task)

- **Rewards**: Box on target (+1.0), All boxes placed (+10.0), Format correct (+0.5), Step penalty (-0.1)

- **Hyperparameters**: `γ_turn`=0.95, `γ_token`=1.0, KL penalty=0.001, Actor LR=1e-6, Critic LR=1e-5

## News

----

**[2025/05]** We've introduced a new modular design for environments and services in VAGEN:

- Enhanced environment framework for easier creation of custom environments

- New service architecture for efficient distributed training

- Check out our new guides:

  - [Creating Environments](docs/creating_environments.md): Learn how to build custom environments

  - [Service Architecture](docs/service_architecture.md): Understand our scalable training infrastructure

## Installation

cd VAGEN

bash scripts/install.sh

```

## Examples

```

# Login to wandb

wandb login

## Reproduce Experiments

# Then, you can run different environments and algorithms:

```bash

# To reproduce our reults, please go to release branch of verl and v25.3.25 of vagen

cd ../verl

git checkout release

cd ../VAGEN

git checkout v25.3.25

wandb login # login into wandb

# Then, you can run

bash vagen/examples/release_experiments/gae.sh # rico-gae

bash vagen/examples/release_experiments/grpo_mask_loss.sh # rico-grpo + loss mask

bash vagen/examples/release_experiments/grpo.sh # rico-grpo

bash vagen/examples/release_experiments/mask_gae_mask_loss_bi_level.sh # trico - turn reward

bash vagen/examples/release_experiments/mask_gae_mask_loss_turnwise_gae.sh # trico - turn reward - bi-level gae + turn-level gae

bash vagen/examples/release_experiments/mask_gae_mask_loss_turnwise_reward_bi_level.sh # trico

bash vagen/examples/release_experiments/mask_gae_mask_loss.sh # aico

bash vagen/examples/release_experiments/mask_gae.sh # aico - loss mask

bash vagen/examples/release_experiments/mask_loss.sh # aico - gae mask

# Frozen Lake Environment

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

| Setting           | GRPO | GAE | Bi-Level GAE | Turn-Wise GAE | Masked-GAE |

|-------------------|------|-----|--------------|---------------|------------|

| with_loss_mask    | ✓    | ✓   | ✓            | ✓             | ✓          |

| multi-turn-reward | ✗    | ✓   | ✓            | ✓             | ✓          |

| with_gae_mask     | ✗    | ✗   | ✓            | ✓             | ✓          |

<img width="800" alt="image" src="./public/1.png" />

### Algorithm Options

<img width="800" alt="image" src="./public/2.png" />

- **GRPO**: Whether to use GRPO

  - `algorithm.adv_estimator=grpo`

- **GAE**: Whether to use GAE

  - `algorithm.adv_estimator=gae`

- **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`

<img width="800" alt="image" src="./public/3.png" />

### 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`

- **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).

# Algorithm Configurations

VAGEN supports several advantage estimation algorithms, each with different properties for training VLM agents. 

## Algorithm Quick Reference

#### RICO (Traditional RL)

```python

algorithm.adv_estimator=grpo

rollout_manager.use_loss_mask=True

rollout_manager.use_gae_mask=False

rollout_manager.use_multi_turn_reward=False

```

#### AICO (Action-centric Optimization)

```python

algorithm.adv_estimator=masked_gae

rollout_manager.use_loss_mask=True

rollout_manager.use_gae_mask=True

rollout_manager.use_multi_turn_reward=False

```

#### TRICO (Turn-aware Optimization)

```python

algorithm.adv_estimator=bi_level_gae

algorithm.high_level_gamma=0.95

rollout_manager.use_loss_mask=True

rollout_manager.use_gae_mask=True

rollout_manager.use_multi_turn_reward=True

```

## Algorithm Settings

The table below summarizes which features are enabled by default with each algorithm:

| Setting           | GRPO | GAE | Bi-Level GAE | Turn-Wise GAE | Masked-GAE |

|-------------------|------|-----|--------------|---------------|------------|

| with_loss_mask    | ✓    | ✓   | ✓            | ✓             | ✓          |

| multi-turn-reward | ✗    | ✓   | ✓            | ✓             | ✓          |

| with_gae_mask     | ✗    | ✗   | ✓            | ✓             | ✓          |

### Algorithm Options

- **GRPO**: Whether to use GRPO

    - `algorithm.adv_estimator=grpo`

- **GAE**: Whether to use GAE

    - `algorithm.adv_estimator=gae`

- **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`

### Algorithm 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`

- **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`

 No newline at end of file

# How to Create New Services

# How to Create New Environments

> NOTICE: Once you've implemented your environment following this guide, VAGEN can be used directly, with the service layer being optional for training acceleration.

This guide explains how to create new environment services for VAGEN's service-based architecture. The service architecture provides a standardized way to manage multiple environments for VLM agent training with enhanced scalability and flexibility.

This guide explains how to create new environments for VAGEN's architecture. Creating custom environments is the foundation for building specialized VLM agent training scenarios. 

## Service Architecture Overview

## Environment Structure Overview

VAGEN uses a client-server architecture for environment management:

VAGEN uses an object-oriented approach for environment management:

- 'BaseEnv': Abstract base class that defines the interface all environments must implement

- 'BaseEnvConfig': Configuration class for environment parameters

- 'Environment'-specific implementations (e.g., SvgEnv)

- `BaseService`: Abstract base class that defines the interface all services must implement

- `BatchEnvClient`: Client that communicates with environment servers (fixed)

- `BatchEnvServer`: Server implementation for hosting environments (fixed)

This architecture enables efficient parallel processing across distributed systems and seamless integration of both rule-based rewards and reward models.

This architecture enables standardized interaction patterns while allowing for customization across different domains and tasks.

## 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 serialization from service to client

```

## Component Hierarchy

|   ├── create_dataset.py         # Store train/test data configs, not real data

│   ├── base/

│   │   ├── base_env.py           # Abstract base class defining the 

│   │   └── base_env_config.py    # Base configuration class for environments

│   ├── [your_env]/               # Your environment implementation

│       ├── env.py                # Your environment class

│       ├── env_config.py         # Your environment configuration

│       └── data/                 # Environment-specific resources (Optional)

|   

├── examples/

|   ├── [your_env]/

│       ├── env_config.yaml       # Your data&env config for create_dataset.py

│       ├── run.sh                # Script for training

```

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**

```

## Creating a New Service Step by Step

## Creating a New Environment Step by Step

### Step 1: Inherit from BaseService

### Step 1: Create Environment Configuration

Create a new class that inherits from `BaseService`. This class must implement all required methods for interacting with environments:

Create a new class that inherits from BaseEnvConfig. This class will define all parameters specific to your environment in `create_dataset.py` by combining your unique requirements in `env_config.yaml` and default requirements in `env_config.py`:

```python

from vagen.env.base_service import BaseService

class MyNewService(BaseService):

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

        # Initialize your environment-specific components here

    def create_environments_batch(self, env_ids, **kwargs):

        # Initialize batch of environments

        # Return initialization status

    def reset_batch(self, env_ids, **kwargs):

        # Reset environments and return observations

    def step_batch(self, env_ids, actions, **kwargs):

        # Process actions and return (observations, dones)

    def compute_reward_batch(self, env_ids, **kwargs):

        # Calculate rewards for each environment

        # Return rewards and any additional info

@dataclass

class MyNewEnvConfig(BaseEnvConfig):

    """Configuration for My New Environment"""

    dataset_name: str = "path/to/dataset"

    data_dir: str = "vagen/env/my_new_env/data"

    seed: int = 42

    # Add your environment-specific parameters here

    def config_id(self) -> str:

        """Generate a unique identifier for this configuration"""

        return f"MyNewEnvConfig(dataset={self.dataset_name},seed={self.seed})"

```

    def get_system_prompts_batch(self, env_ids, **kwargs):

        # Return system prompts for each environment

### Step 2: Implement Environment Class

    def close_batch(self, env_ids, **kwargs):

        # Clean up resources for environments

```

### Step 2: Observation Serialization

When passing observations between the service and client, ensure proper serialization:

Create a new class that inherits from BaseEnv. This class must implement all required methods:

```python

from vagen.env.base.base_env import BaseEnv

from typing import Dict, Tuple

import random

class MyNewEnv(BaseEnv):

    def __init__(self, config):

        self.config = config

        self.done = False

    def step(self, llm_raw_response) -> Tuple[Dict, float, bool, Dict]:

        """Process an action from the LLM and return the next state"""

        parsed_action = parse_llm_raw_response(llm_raw_response)

        action_valid = parsed_action['is_valid']

        action_effective = action_valid  # Simplification for example

        # Update environment state based on action

        obs = {

            'obs_str': "Observation after action",

            'multi_modal_data': {}  # Add any images or audio here

        }

        reward = 0.0 if not action_valid else 0.5

        self.done = False  # Update based on task completion

        info = {

            "metrics": {

                'success': False,

                'action_is_effective': action_effective,

                'action_is_valid': action_valid,

            },

            "llm_raw_response": llm_raw_response,

            "llm_response": parsed_action,

        }

        return obs, reward, self.done, info

    def reset(self, seed=None) -> Tuple[Dict, Dict]:

        """Reset the environment to initial state"""

        if seed is not None:

            random.seed(seed)

        self.done = False

        obs = {

            'obs_str': "Initial observation text",

            'multi_modal_data': {}

        }

        return obs, info #(Optional, could be empty)

    def system_prompt(self) -> str:

        """Define the system prompt for the LLM"""

        return "You are an agent in the MyNewEnv environment. Your goal is to [describe task]."

    def compute_reward(self) -> float:

        """Calculate final episode reward"""

        return 0.0  # Calculate based on task completion

    def close(self):

        """Clean up any resources"""

        pass

```

from vagen.utils.serial import serialize_observation, deserialize_observation

### Step 3: Make Sure Input/Output Format Details

# On the service side

## at the end of reset_batch()

serialized_obs = serialize_observation(original_observation)

## at the end of step_batch()

serialized_step = serialize_step_result(observation, reward, done, info)

#### Environment Observations

Step Observations must follow this structure:

```python

{

    'obs_str': "Text observation with <image> or <audio> placeholders",

    'multi_modal_data': {

        '<image>': [image_data_1, image_data_2, ...],

        '<audio>': [audio_data_1, audio_data_2, ...],

    }

}

```

**Notice**: number of `image_place_holder(<image>)` in `obs_str` must match with number of `image_data` in `multi_modal_data`

### Step 3: Register your Environment

Register your env in `env/__init__.py`

```

from vagen.env.NEW_service import MyNewService

# Register your service

REGISTERED_ENV["my_new_env"] = MyNewService

#### Environment Info Dictionary

The info dictionary provides additional context and metrics:

```python

{

    "metrics": {

        'success': bool,  # Did the agent complete the task?

        'action_is_effective': bool,  # Was the action meaningful?

        'action_is_valid': bool,  # Was the action syntactically correct?

        # Add additional custom metrics

    },

    "llm_raw_response": str,  # Original response from LLM

    "llm_response": dict,  # Parsed response with structured format

}

```

### Step 4: Define your env config and script

> Please refer to `[Configuration](config.md)`

Define your env config and running script in `examples/`

### Step 4: Testing Your Environment

Create a basic script below your `env.py` to test your environment:

```python

# Create environment

config = MyNewEnvConfig()

env = MyNewEnv(config)

# Reset environment

obs, info = env.reset(seed=42)

print("Initial observation:", obs['obs_str'])

# Test step with mock LLM response

mock_llm_response = "Action1, Action2, Action3"

next_obs, reward, done, info = env.step(mock_llm_response)

print("Next observation:", next_obs['obs_str'])

print("Reward:", reward)

print("Done:", done)

print("Action valid:", info['metrics']['action_is_valid'])

print("Action effective:", info['metrics']['action_is_effective'])

# Clean up

env.close()

```

### Step 5: Integration with Service Layer (Optional)

Please refer to `Frozenlake/service.py` for better service structure understanding

 No newline at end of file

For training acceleration and distributed processing, you can integrate your environment with the VAGEN service layer. This step is optional but recommended for large-scale training. See the "[Create your Own Service](create-service.md)" section for details.