ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s

This commit is contained in:
2026-08-20 13:12:50 +00:00
commit b119135836
10275 changed files with 3284984 additions and 0 deletions
@@ -0,0 +1,246 @@
<div align="center">
# AWorld Train
*Framework-agnostic training adapters, examples, and utilities for training AWorld agents with external RL/training frameworks*
[![License: MIT][license-image]][license-url]
</div>
---
AWorld Train provides a bridge between the AWorld agent ecosystem and various external training frameworks like Reinforcement Learning (RL) libraries. It is designed to be framework-agnostic, allowing you to bring your AWorld agents to your favorite training environments.
The following diagram illustrates the overall architecture and the interaction between the Environment host and Training cluster:
![Architecture Diagram](../readme_assets/train_env_agent_architecture.png)
## Environment host construction
First, you need to set up the environment where the agent's tools will run.
Choose a machine (which can be a training machine).
Machine sizing recommendation:
- For capacity planning, allocate roughly **2C4G** per concurrent worker.
- Example: for concurrency=8, plan for **~16C and ~32G**.
```bash
# git clone AWorld
git clone git@github.com:inclusionAI/AWorld.git
cd /path/to/AWorld
cp ./env/gaia-mcp-server/mcp_servers/.env_template ./env/gaia-mcp-server/mcp_servers/.env
```
Edit ./env/gaia-mcp-server/mcp_servers/.env to configure authentication tokens for any required tools.
```.env
JINA_API_KEY=<YOUR_JINA_API_KEY>
TAVILY_API_KEY=<YOUR_TAVILY_API_KEY>
GOOGLE_API_KEY=<YOUR_GOOGLE_API_KEY>
GOOGLE_CSE_ID=<YOUR_GOOGLE_CSE_ID>
DATALAB_API_KEY=<YOUR_DATALAB_API_KEY>
E2B_API_KEY=<YOUR_E2B_API_KEY>
MCP_LLM_BASE_URL=<YOUR_MCP_LLM_BASE_URL>
MCP_LLM_MODEL_NAME=<YOUR_MCP_LLM_MODEL_NAME>
MCP_LLM_API_KEY=<YOUR_MCP_LLM_API_KEY>
BROWSERUSE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
BROWSERUSE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
BROWSERUSE_LLM_API_KEY=${MCP_LLM_API_KEY}
CODE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
CODE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
CODE_LLM_API_KEY=${MCP_LLM_API_KEY}
THINK_LLM_BASE_URL=${MCP_LLM_BASE_URL}
THINK_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
THINK_LLM_API_KEY=${MCP_LLM_API_KEY}
GUARD_LLM_BASE_URL=${MCP_LLM_BASE_URL}
GUARD_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
GUARD_LLM_API_KEY=${MCP_LLM_API_KEY}
AUDIO_LLM_BASE_URL=${MCP_LLM_BASE_URL}
AUDIO_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
AUDIO_LLM_API_KEY=${MCP_LLM_API_KEY}
IMAGE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
IMAGE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
IMAGE_LLM_API_KEY=${MCP_LLM_API_KEY}
VIDEO_LLM_BASE_URL=${MCP_LLM_BASE_URL}
VIDEO_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
VIDEO_LLM_API_KEY=${MCP_LLM_API_KEY}
```
Next, run the startup script to launch the MCP server locally:
```bash
cd /path/to/Aworld
# use --docker_dir to specify the docker directory to build
# e.g., --docker_dir=gaia-mcp-server
python -m env.train_env --docker_dir=gaia-mcp-server
```
Once the MCP server starts successfully, it will output the connection details:
```bash
{
"ip": "1xx.1xx.x.xx",
"port": 8000,
"token": "eyJhbGciOi...rYmQ"
}
```
You will need the ip, port and token from this output for the next step, where you'll configure the Agent on your training machine.
For instructions on deploying the environment on Kubernetes, please refer to [`../env/README.md`](../env/README.md).
## Training cluster Setup
### 1. Create an Agent or Swarm
Now, on the training cluster machine, you must make the MCP service credentials available to your agent. Use the ip, port and token from the [Environment host](#environment-host) section and export them as environment variables or add them to a `.env` file:
```bash
# export them as environment variables
# replace <ip>, <port> and <token> with the ip, port and token from Step 1
export MCP_SERVER_URL=http://<ip>:<port>/mcp
export MCP_SERVER_TOKEN=<token>
# or add them to `.env` file
# echo "MCP_SERVER_URL=http://<ip>:<port>/mcp" >> .env
# echo "MCP_SERVER_TOKEN=<token>" >> .env
```
Then install aworld and RL framework:
```bash
# Python>=3.10 is recommended.
# Install AWorld
pip install aworld
# Framework-specific deps (VeRL example)
pip install verl==0.5.0
```
With the connection details configured, you can define your agent within your chosen training framework. For VeRL, this is accomplished by implementing a custom `AgentLoop`.
For example, `GaiaAgentLoop` inherits from `AworldAgentLoop` and implements the `build_agents` method.
```python
from aworld.agents.llm_agent import Agent
from aworld.config import AgentConfig
from train.adapter.verl.aworld_agent_loop import AworldAgentLoop
from train.adapter.verl.common import get_agent_tool_env_and_servers
class GaiaAgentLoop(AworldAgentLoop):
def build_agents(self):
# Get the environment configuration and server details.
# Note: The MCP server must be running (Step 1) and the
# MCP_SERVER_URL/MCP_SERVER_TOKEN environment variables must be set.
gaia_env_config, gaia_env_servers = get_agent_tool_env_and_servers()
return Agent(
conf=AgentConfig(
# Get the dynamic llm server address from the server manager.
# The llm server is launched within VeRL.
llm_base_url=self.get_llm_server_address(),
llm_model_name=self.get_llm_server_model_name(),
),
name="gaia_super_agent",
system_prompt="YOUR SYSTEM PROMPT",
# MCP tool configuration for the agent
mcp_config=gaia_env_config,
mcp_servers=gaia_env_servers,
)
```
### 2. Run Training
Before run training, specify your custom `AgentLoop` in the `agent.yaml`:
```yaml
# In agent.yaml
- name: gaia_agent
_target_: train.examples.train_gaia_with_aworld_verl.custom_agent_loop.GaiaAgentLoop
```
Finally, run the training script. This script is typically a `run.sh` file based on the VeRL example.
```bash
bash run.sh
```
This script handles the training loop, reward calculation, and agent updates, orchestrated by VeRL.
Please refer to the [VeRL documentation](https://verl.readthedocs.io/en/latest/examples/config.html) for parameter settings in `run.sh`.
A complete, runnable example, including a `run.sh` script tailored for `GaiaAgentLoop`, is available in [`./examples/train_gaia_with_aworld_verl/`](./examples/train_gaia_with_aworld_verl/).
## Advanced Tutorial
### How to Create a Complex Swarm
Instead of a single agent, you can also train a multi-agent swarm. Simply have your `build_agents` method (or equivalent setup function) return a `Swarm` object instead of a single `Agent`. AWorld and the training adapter will handle the rest.
```python
# In your AgentLoop or setup file
def build_agents(self, ...) -> Union[Agent, Swarm]:
# ... (create individual agents)
agent_to_be_train = Agent(
conf=AgentConfig(
# For the agent to be trained, llm_base_url and llm_model_name are obtained from the services launched by VeRL
llm_base_url=self.get_llm_server_address(),
llm_model_name=self.get_llm_server_model_name(),
),
)
plan_agent = Agent(
conf=AgentConfig(
# Provide a ready-to-use OpenAI-compatible llm service address, model name, and api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
exe_agent = Agent(
conf=AgentConfig(
# Provide a ready-to-use OpenAI-compatible llm service address, model name, and api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
sum_agent = Agent(
conf=AgentConfig(
# Provide a ready-to-use OpenAI-compatible llm service address, model name, and api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
# Return a Swarm composed of your agents
return Swarm(
agent_to_be_train, plan_agent, exe_agent, sum_agent,
# ... other swarm configuration
)
```
### How to Integrate with Other Training Frameworks
AWorld Train is designed for extensibility. To add support for a new training framework (e.g., "Swift"), you would typically need to:
1. **Create a new Adapter**: Inside the `train/adapter/` directory, create a new folder for your framework (e.g., `swift/`).
2. **Implement the Core Logic**: Create a primary class (e.g., `AworldAgentTrainer`) that inherits from a base class of the target framework. This class will be responsible for:
* Receiving tasks or observations from the framework's environment.
* Run the AWorld agent (`Runners.sync_run(input=input, agent=agent)`) to get an action.
* Returning the agent's response back to the framework.
* Handling rewards and updates.
3. **Create an Example**: Add a new example in the `train/examples/` directory to demonstrate how to use the new adapter.
You can refer to the existing `verl` adapter (`train/adapter/verl/`) as a reference implementation.
---
<div align="center">
**AWorld Train** — Bring your AWorld agents to your favorite training frameworks
[license-image]: https://img.shields.io/badge/License-MIT-yellow.svg
[license-url]: https://opensource.org/licenses/MIT
</div>
@@ -0,0 +1,248 @@
<div align="center">
# AWorld Train
*为使用 AWorld 构建的智能体,提供与外部 RL/训练框架对接的、与框架无关的适配层、可运行示例与通用工具*
[![License: MIT][license-image]][license-url]
</div>
---
AWorld Train 为 AWorld 智能体生态系统和各种外部训练框架(如强化学习库)之间提供了一座桥梁。它被设计为框架无关的,可以选择你喜欢的训练环境使用AWorld 智能体。
下图说明了环境和训练集群之间的整体架构和交互:
![架构示意图](../readme_assets/train_env_agent_architecture.png)
## 环境构建
首先,您需要设置智能体工具将要运行的环境。
选择一台机器(也可以是训练机)。
机器规格建议:
- 为进行容量规划,为每个并发工作进程分配大约 **2C4G**
- 示例:对于8个并发,计划需要 **约16C32G**
```bash
# 克隆 AWorld 仓库
git clone git@github.com:inclusionAI/AWorld.git
cd /path/to/AWorld
cp ./env/gaia-mcp-server/mcp_servers/.env_template ./env/gaia-mcp-server/mcp_servers/.env
```
编辑 `./env/gaia-mcp-server/mcp_servers/.env` 以配置任何需要身份验证的工具的令牌。
```.env
JINA_API_KEY=<YOUR_JINA_API_KEY>
TAVILY_API_KEY=<YOUR_TAVILY_API_KEY>
GOOGLE_API_KEY=<YOUR_GOOGLE_API_KEY>
GOOGLE_CSE_ID=<YOUR_GOOGLE_CSE_ID>
DATALAB_API_KEY=<YOUR_DATALAB_API_KEY>
E2B_API_KEY=<YOUR_E2B_API_KEY>
MCP_LLM_BASE_URL=<YOUR_MCP_LLM_BASE_URL>
MCP_LLM_MODEL_NAME=<YOUR_MCP_LLM_MODEL_NAME>
MCP_LLM_API_KEY=<YOUR_MCP_LLM_API_KEY>
BROWSERUSE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
BROWSERUSE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
BROWSERUSE_LLM_API_KEY=${MCP_LLM_API_KEY}
CODE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
CODE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
CODE_LLM_API_KEY=${MCP_LLM_API_KEY}
THINK_LLM_BASE_URL=${MCP_LLM_BASE_URL}
THINK_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
THINK_LLM_API_KEY=${MCP_LLM_API_KEY}
GUARD_LLM_BASE_URL=${MCP_LLM_BASE_URL}
GUARD_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
GUARD_LLM_API_KEY=${MCP_LLM_API_KEY}
AUDIO_LLM_BASE_URL=${MCP_LLM_BASE_URL}
AUDIO_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
AUDIO_LLM_API_KEY=${MCP_LLM_API_KEY}
IMAGE_LLM_BASE_URL=${MCP_LLM_BASE_URL}
IMAGE_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
IMAGE_LLM_API_KEY=${MCP_LLM_API_KEY}
VIDEO_LLM_BASE_URL=${MCP_LLM_BASE_URL}
VIDEO_LLM_MODEL_NAME=${MCP_LLM_MODEL_NAME}
VIDEO_LLM_API_KEY=${MCP_LLM_API_KEY}
```
接下来,运行启动脚本以在本地启动 MCP 服务器:
```bash
cd /path/to/Aworld
# --docker_dir 参数指定需要构建的env对应docker目录
# e.g., --docker_dir=gaia-mcp-server
python -m env.train_env --docker_dir=gaia-mcp-server
```
MCP 服务器成功启动后,将输出连接详细信息:
```bash
{
"ip": "1xx.1xx.x.xx",
"port": 8000,
"token": "eyJhbGciOi...rYmQ"
}
```
您将需要此输出中的 `ip`、`port` 和 `token`,用于下一步在训练机上配置智能体。
有关在 Kubernetes 上部署环境的说明,请参阅 [`../env/README.md`](../env/README.md)。
## 训练集群设置
### 1. 创建智能体或智能体集群
现在,在训练集群机器上,您必须使 MCP 服务凭据对您的智能体可用。使用[环境构建](#环境构建)部分中的 `ip`、`port` 和 `token`,并将它们导出为环境变量或添加到 `.env` 文件中:
```bash
# 导出为环境变量
# 将 <ip>、<port> 和 <token> 替换为环境构建中的 ip、port 和 token
export MCP_SERVER_URL=http://<ip>:<port>/mcp
export MCP_SERVER_TOKEN=<token>
# 或将它们添加到 .env 文件中
# echo "MCP_SERVER_URL=http://<ip>:<port>/mcp" >> .env
# echo "MCP_SERVER_TOKEN=<token>" >> .env
```
然后安装 `aworld` 和强化学习框架:
```bash
# 推荐使用 Python>=3.10。
# 安装 AWorld
pip install aworld
# 安装特定框架的依赖(以 VeRL 为例)
pip install verl==0.5.0
```
配置好连接详细信息后,您可以在所选的训练框架内定义您的智能体。对于 VeRL,这是通过实现一个自定义的 `AgentLoop` 来完成的。
例如,`GaiaAgentLoop` 继承自 `AworldAgentLoop` 并实现了 `build_agents` 方法。
```python
from aworld.agents.llm_agent import Agent
from aworld.config import AgentConfig
from train.adapter.verl.aworld_agent_loop import AworldAgentLoop
from train.adapter.verl.common import get_agent_tool_env_and_servers
class GaiaAgentLoop(AworldAgentLoop):
def build_agents(self):
# 获取环境配置和服务器详细信息。
# 注意:MCP 服务器必须正在运行(环境构建),并且
# MCP_SERVER_URL/MCP_SERVER_TOKEN 环境变量必须已设置。
gaia_env_config, gaia_env_servers = get_agent_tool_env_and_servers()
return Agent(
conf=AgentConfig(
# 从服务管理器获取动态的 llm 服务地址。
# llm 服务是在 VeRL 中启动的。
llm_base_url=self.get_llm_server_address(),
llm_model_name=self.get_llm_server_model_name(),
),
name="gaia_super_agent",
system_prompt="你的系统提示",
# 智能体的 MCP 工具配置
mcp_config=gaia_env_config,
mcp_servers=gaia_env_servers,
)
```
### 2. 运行训练
在运行训练之前,请在 `agent.yaml` 中指定您的自定义 `AgentLoop`
```yaml
# 在 agent.yaml 中
- name: gaia_agent
_target_: train.examples.train_gaia_with_aworld_verl.custom_agent_loop.GaiaAgentLoop
```
最后,运行训练脚本。该脚本通常是基于 VeRL 示例的 `run.sh` 文件。
```bash
bash run.sh
```
此脚本处理由 VeRL 编排的AgentLoop、奖励计算函数和训练流程。
有关 `run.sh` 中的参数设置,请参阅 [VeRL 文档](https://verl.readthedocs.io/en/latest/examples/config.html)。
一个完整的、可运行的示例,包括为 `GaiaAgentLoop` 定制的 `run.sh` 脚本,可在 [`./examples/train_gaia_with_aworld_verl/`](./examples/train_gaia_with_aworld_verl/) 中找到。
## 进阶教程
### 如何创建复杂的多智能体集群 (Swarm)
除了单个智能体,您还可以训练一个多智能体集群。只需让您的 `build_agents` 方法(或等效的设置函数)返回一个 `Swarm` 对象而不是单个 `Agent` 对象即可。AWorld 和训练适配器将处理剩下的部分。
```python
# 在自定义的AgentLoop中
def build_agents(self, ...) -> Union[Agent, Swarm]:
# ... 创建多个agent
agent_to_be_train = Agent(
conf=AgentConfig(
# 对于要训练的agentllm_base_url 和 llm_model_name 是从 VeRL 启动的服务中获取的
llm_base_url=self.get_llm_server_address(),
llm_model_name=self.get_llm_server_model_name(),
),
)
plan_agent = Agent(
conf=AgentConfig(
# 提供一个即用型的 OpenAI 兼容的 llm 服务地址、模型名称和 api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
exe_agent = Agent(
conf=AgentConfig(
# 提供一个即用型的 OpenAI 兼容的 llm 服务地址、模型名称和 api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
sum_agent = Agent(
conf=AgentConfig(
# 提供一个即用型的 OpenAI 兼容的 llm 服务地址、模型名称和 api_key
llm_base_url="",
llm_model_name="",
llm_api_key=""
),
)
# 返回由以上定义的智能体组成的Swarm
return Swarm(
agent_to_be_train, plan_agent, exe_agent, sum_agent,
# ... 其他Swarm配置
)
```
### 如何集成其他训练框架
AWorld Train 被设计为可扩展的。要为新的训练框架(例如 “Swift”)添加支持,通常需要:
1. **创建新的适配器**:在 `train/adapter/` 目录内,为您的框架创建一个新文件夹(例如 `swift/`)。
2. **实现核心逻辑**:创建一个主类(例如 `AworldAgentTrainer`),它继承自目标框架的某个基类。这个类将负责:
* 从框架的环境中接收任务或观察结果。
* 运行 AWorld 智能体(`Runners.sync_run(input=input, agent=agent)`)以获取动作。
* 将智能体的响应返回给框架。
* 处理奖励和更新。
3. **创建示例**:在 `train/examples/` 目录中添加一个新示例,以演示如何使用新的适配器。
可以参考现有的 `verl` 适配器(`train/adapter/verl/`)作为参考实现。
---
<div align="center">
**AWorld Train** — 让你的 AWorld 智能体快速接入主流训练框架
[license-image]: https://img.shields.io/badge/License-MIT-yellow.svg
[license-url]: https://opensource.org/licenses/MIT
</div>
@@ -0,0 +1,2 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
@@ -0,0 +1,284 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import abc
import asyncio
import json
import re
from typing import List, Optional, Dict, Any, Union
from aworld.agents.llm_agent import Agent
from aworld.core.agent.swarm import Swarm
from aworld.core.task import TaskResponse
from aworld.runner import Runners
from aworld.utils.common import sync_exec
from swift.llm import RequestConfig
from swift.llm.infer.protocol import ChatCompletionResponse
from swift.trainers.rlhf_trainer.grpo_trainer import InputsType, GRPOTrainer, logger
from transformers import AutoTokenizer
from trl.extras.profiling import profiling_context
class AworldTrainer(GRPOTrainer):
def _engine_infer(
self,
infer_requests: InputsType,
request_config: Optional[RequestConfig] = None,
*,
use_tqdm: Optional[bool] = False,
) -> List[ChatCompletionResponse]:
with profiling_context(self, 'generate'):
if self.vllm_mode != 'server':
return self.engine.infer(infer_requests, request_config, use_tqdm=use_tqdm)
request_keys = ['messages', 'images', 'audios', 'videos', 'tools', 'objects']
infer_requests = [{
**{k: request[k]
for k in request_keys if k in request},
**({
'data_dict': {k: request[k]
for k in request if k not in request_keys}
} if self.multi_turn_scheduler and self.vllm_use_async_engine else {})
} for request in infer_requests]
self._process_infer_requests_images(infer_requests)
return self.run_infer(infer_requests)
def run_infer(self, infer_requests: List[Dict[str, Any]]) -> List[ChatCompletionResponse]:
workers = [asyncio.create_task(self._rollout(req)) for req in infer_requests]
results = sync_exec(asyncio.gather, *workers)
return self.convert_agent_output(results, infer_requests)
async def _rollout(self, req: Dict[str, Any]):
agent = self.build_agents()
result = await self.run_agents(req, agent)
return result
@abc.abstractmethod
def build_agents(self) -> Union[Agent, Swarm]:
"""Build single- or multi-agent"""
async def run_agents(self, input, agent):
# collect trajectory
if isinstance(agent, Swarm):
result = Runners.sync_run(input=input, swarm=agent)
else:
result = Runners.sync_run(input=input, agent=agent)
return result
def convert_agent_output(self,
results: List[TaskResponse],
infer_requests: List[Dict[str, Any]]) -> List[ChatCompletionResponse]:
message_final_merge = []
for result in results:
trajectory = result.trajectory
last_exp_data = trajectory[-1]['exp_data']
task_id = trajectory[0]['exp_meta']['task_id'].split('_')[1]
message_final = []
message = last_exp_data["messages"]
answer_flag = 0
for i in range(len(message)):
actions = last_exp_data.get('actions', [])
if actions:
actions_str = json.dumps(actions)
if '<answer>' in actions_str and '</answer>' in actions_str:
match = re.search(r'<answer>(.*?)</answer>', actions_str, re.DOTALL)
if match:
answer_flag = 1
logger.info(f"{task_id} answer content: {match.group(1)}")
else:
logger.warning(f"{task_id} no answer content found.")
if message[i]["role"] in ["system", "user"]:
message_final.append(
{
"role": message[i]["role"],
"content": message[i]["content"],
}
)
elif message[i]["role"] == "assistant" and "tool_calls" in message[i].keys():
if message[i]["tool_calls"][0]["function"]["arguments"]:
arguments = json.loads(message[i]["tool_calls"][0]["function"]["arguments"])
else:
arguments = ""
function_call = {
"name": message[i]["tool_calls"][0]["function"]["name"],
"arguments": arguments
}
if message[i]["content"] != "" and message[i]["content"] is not None:
message_final.append(
{
"role": "assistant",
"content": message[i]["content"],
}
)
message_final.append(
{
"role": "tool_call",
"content": json.dumps(function_call, ensure_ascii=False),
}
)
elif message[i]["role"] == "tool":
last_content = message[i - 1]["content"]
if last_content is None:
last_content = ""
message_final.append(
{
"role": "tool",
"content": message[i]["content"].replace(last_content, ""),
}
)
else:
logger.warning(f"Unknown message role: {message[i]['role']}")
tokenizer = AutoTokenizer.from_pretrained(self.args.model_init_kwargs)
try:
response = last_exp_data["actions"][0]["policy_info"]
if response:
message_final.append(
{
"role": "assistant",
"content": response
}
)
else:
message_final.append(
{
"role": "assistant",
"content": "No response was received. Please try again later."
}
)
message_final = truncate_messages_fast(message_final, tokenizer)
status = "success" if answer_flag == 1 else "length"
message_final_merge.append((message_final, status, task_id))
except:
message_final.append({
"role": "assistant",
"content": "No response was received. Please try again later."
})
message_final = truncate_messages_fast(message_final, tokenizer)
message_final_merge.append((message_final, "length", task_id))
return self.pad_list_to_length(message_final_merge, infer_requests)
def pad_list_to_length(self, message_final_merge, infer_requests) -> List[ChatCompletionResponse]:
unique_task_ids = []
for msg in message_final_merge:
task_id = msg[2]
if task_id not in unique_task_ids:
unique_task_ids.append(task_id)
# Group by task_id
task_groups = {task_id: [] for task_id in unique_task_ids}
for item in message_final_merge:
messages, status, task_id = item
if task_id in task_groups:
task_groups[task_id].append(item)
# Ensure each group has exactly num_generations samples
for task_id in unique_task_ids:
# If this task_id has no samples, construct fallback data
if len(task_groups[task_id]) == 0:
for _infer_request in infer_requests:
# Check if the first message content matches the task_id
if task_id == _infer_request["messages"][0]["content"]:
fallback_completion = {
"role": "assistant",
"content": "No response was received. Please try again later."
}
new_messages = _infer_request["messages"].copy()[1:]
new_messages.append(fallback_completion)
task_groups[task_id].append((new_messages, "length", task_id))
break
# # Ensure we have exactly num_generations samples
# while len(task_groups[task_id]) < num_generations/len(unique_task_ids):
# success_samples = [item for item in task_groups[task_id] if item[1] == "success"]
# if success_samples:
# task_groups[task_id].append(random.choice(success_samples))
# else:
# task_groups[task_id].append(random.choice(task_groups[task_id]))
num_generations = len(infer_requests)
current_count = len(task_groups[task_id])
if current_count >= num_generations / len(unique_task_ids):
continue
# Get success samples if available, otherwise use all samples
success_samples = [item for item in task_groups[task_id] if item[1] == "success"]
samples_to_cycle = success_samples if success_samples else task_groups[task_id]
# Calculate how many more we need
needed = int(num_generations / len(unique_task_ids)) - current_count
# Add samples in a cycling manner
for i in range(int(needed)):
task_groups[task_id].append(samples_to_cycle[i % len(samples_to_cycle)])
# Combine all groups and convert back to 2-tuples for final output
final_result = []
for task_id in unique_task_ids:
for item in task_groups[task_id]:
messages, status, _ = item
final_result.append((messages, status))
return final_result
def truncate_messages_fast(
messages: List[Dict[str, Any]],
tokenizer: Any,
max_length: int = 131072,
tools: Optional[List] = None
) -> List[Dict[str, Any]]:
"""Simplifies message list truncation by removing entire messages from the end
to fit within max_length, with a final role check.
Core Logic:
1. First, removes messages from the end of the list one by one until the
total token count is within `max_length`.
2. After ensuring the length is acceptable, it performs a final check on the
last remaining message.
3. If the last message's role is not 'assistant' or 'tool_call', it is
also removed. This check is repeated until the last message has a valid
role or the list becomes empty.
4. This function does not partially truncate message content.
Args:
messages (List[Dict[str, Any]]): A list of message dictionaries.
tokenizer: The tokenizer instance to calculate token count.
max_length (int, optional): The target maximum number of tokens. Defaults to 131072.
tools (Optional[List], optional): A list of tools that might be needed when applying
the chat template. Defaults to None.
Returns:
List[Dict[str, Any]]: The truncated list of messages.
"""
truncated_messages = list(messages)
def get_current_tokens(msgs: List[Dict[str, Any]]) -> int:
if not msgs:
return 0
# The return value of apply_chat_template can be a list of token IDs or a string
# We use len() to get the count, which works for both cases.
return len(tokenizer.apply_chat_template(msgs, tools=tools, add_generation_prompt=False))
# 1. Truncate from the end based on length
# The `and truncated_messages` ensures we don't pop from an empty list
while get_current_tokens(truncated_messages) > max_length and truncated_messages:
truncated_messages.pop() # pop() removes the last item
# 2. Ensure the last remaining message has a valid role ('assistant' or 'tool_call')
# This loop handles cases where multiple invalid messages are at the end (e.g., ..., tool, user)
while truncated_messages:
last_message_role = truncated_messages[-1].get("role")
if last_message_role in ('assistant', 'tool_call'):
# The last message is valid, so we are done.
break
else:
# The last message is not of the required role, remove it and check again.
truncated_messages.pop()
return truncated_messages
@@ -0,0 +1,27 @@
# VERL Adapter (AWorld Train)
This module hosts the VERL integration for AWorld training workflows.
- aworld_agent_loop.py: Base class bridging VERL AgentLoop with AWorld agents.
- common.py:
- Utilities for converting trajectories/messages to VERL AgentLoopOutput.
- Utilities for getting MCP server configuration.
## Usage
Import adapter entrypoints from your example code:
```python
from train.adapter.verl.aworld_agent_loop import AworldAgentLoop
```
Then implement your example-specific loop:
```python
class MyLoop(AworldAgentLoop):
def build_agents(self):
...
```
## Adding New Features
- Avoid putting example-specific code here; that belongs in train/examples/.
## Notes
- Prefer small, composable utilities and explicit public APIs.
@@ -0,0 +1,179 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import abc
import json
import logging
import os
import uuid
from typing import Any, List, Dict, Union
from aworld.agents.llm_agent import Agent
from aworld.config.agent_loader import _load_yaml
from aworld.core.agent.swarm import Swarm
from aworld.runner import Runners
from aworld.logs.util import logger
from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput
from train.adapter.verl.common import to_agent_loop_output
logger.setLevel(logging.INFO)
logger.propagate = False
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
class AworldAgentLoop(AgentLoopBase):
__metaclass__ = abc.ABCMeta
@abc.abstractmethod
async def build_agents(self) -> Union[Agent, Swarm]:
"""Build single- or multi-agent"""
async def get_llm_server_address(self, server_name: str = None) -> str:
server = self.server_manager._choose_server(server_name or uuid.uuid4().hex)
base_url = await server.get_server_address.remote()
base_url = f"http://{base_url}/v1"
logger.info(f"get_server_address#base_url: {base_url}")
return base_url
async def get_llm_server_model_name(self):
model_name = "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:])
logger.info(f"get_server_model_name#model_name: {model_name}")
return model_name
# main branch
# async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
# messages = list(kwargs["raw_prompt"])
# release 0.5.0
async def run(self, messages: list, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
agent = await self.build_agents()
self.agent = agent
result = await self.run_agents(messages[0], agent)
res = result.trajectory
# build agent loop output
output = await self.convert_agent_output(trajectory=res,
response_length=self.config.actor_rollout_ref.rollout.response_length)
return output
async def run_agents(self, input, agent):
if isinstance(input, dict):
input = input.get("content", "")
# collect trajectory
if isinstance(agent, Swarm):
result = Runners.sync_run(input=input, swarm=agent)
else:
result = Runners.sync_run(input=input, agent=agent)
return result
async def get_agent_tool_config(self, config_path: str) -> Dict[str, Any]:
"""Load tool configuration, preferring YAML with simple fields.
Priority:
1) agent_tools.yaml (simple user config with url, Authorization, MCP_SERVERS)
2) mcp.json (legacy full config)
"""
# 1) Try YAML (simple schema)
try:
import yaml # Local import to avoid hard dependency at import time
if os.path.exists(config_path):
src = _load_yaml(config_path)
url = src.get('url', '')
authorization = src.get('Authorization', '')
mcp_servers_value = src.get('MCP_SERVERS', '')
# Normalize servers to comma-separated string for header and list for internal
if isinstance(mcp_servers_value, list):
mcp_servers_str = ','.join([str(s).strip() for s in mcp_servers_value if str(s).strip()])
else:
mcp_servers_str = str(mcp_servers_value or '').strip()
# Build internal full mcp_config
server_name = src.get('server_name', 'aworld-mcp')
server_type = src.get('type', 'streamable-http')
timeout = src.get('timeout', 600)
sse_read_timeout = src.get('sse_read_timeout', 600)
client_session_timeout_seconds = src.get('client_session_timeout_seconds', 600)
if url:
mcp_config = {
"mcpServers": {
server_name: {
"type": server_type,
"url": url,
"headers": {
"Authorization": authorization,
"MCP_SERVERS": mcp_servers_str,
},
"timeout": timeout,
"sse_read_timeout": sse_read_timeout,
"client_session_timeout_seconds": client_session_timeout_seconds,
}
}
}
return mcp_config
except Exception as err:
print(f"Error loading YAML tool config err: {err}")
# 2) Fallback to legacy JSON
try:
if os.path.exists(config_path):
with open(config_path, "r") as f:
return json.load(f)
except Exception as err:
print(f"Error loading tool config[{config_path}] err is : {err}")
def get_num_turns(self, trajectory: List[Dict[str, Any]]):
return len(trajectory)
async def convert_agent_output(self, trajectory: List[Dict[str, Any]], response_length: int) -> AgentLoopOutput:
"""Convert trajectory to AgentLoopOutput.
Args:
trajectory (List[Dict[str, Any]]): List of agent execution trajectory.
response_length (int): Max length of response.
Returns:
AgentLoopOutput: agent loop output trajectory used for training.
"""
if not trajectory:
raise Exception("Trajectory is empty")
num_turns = self.get_num_turns(trajectory)
messages = trajectory[-1].get("exp_data", {}).get("messages", [])
if not messages:
return AgentLoopOutput(
prompt_ids=[],
response_ids=[],
response_mask=[],
num_turns=num_turns,
metrics={},
)
if messages[-1].get("role") != "assistant":
logger.warning(f"Found last message with role '{messages[-1].get('role')}', but expected 'assistant'. Truncating trailing 'tool' messages.")
last_non_tool_index = -1
for i in range(len(messages) - 1, -1, -1):
if messages[i].get("role") != "tool":
last_non_tool_index = i
break
if last_non_tool_index != -1:
messages = messages[:last_non_tool_index + 1]
else:
messages = []
output = await to_agent_loop_output(tokenizer=self.tokenizer,
messages=messages,
response_length=response_length,
tools=self.agent.tools)
return output
@@ -0,0 +1,202 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import asyncio
import json
import os
from typing import List, Dict, Any
from transformers import AutoTokenizer
from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, AgentLoopMetrics
async def to_agent_loop_output(tokenizer: AutoTokenizer,
messages: List[Dict[str, Any]],
response_length: int,
tools: Dict[str, Any] = None) -> AgentLoopOutput:
"""Convert messages to AgentLoopOutput.
Args:
tokenizer (AutoTokenizer): Tokenizer for tokenize messages.
messages (List[Dict[str, Any]]): List of messages in OpenAI request format.
response_length (int): Max length of response.
tools: Tool list used by the agent.
Returns:
AgentLoopOutput: agent loop output trajectory used for training.
"""
# Ensure tools is iterable for chat templates that iterate over tools
if tools is None:
tools = []
# Normalize messages to satisfy chat templates expectations
def _normalize_message(msg: Dict[str, Any]) -> Dict[str, Any]:
normalized = dict(msg)
# content may be None when assistant only returns tool_calls; make it empty string
if normalized.get("content") is None:
normalized["content"] = ""
# Ensure tool_calls.function.arguments is a string (many templates expect str)
if isinstance(normalized.get("tool_calls"), list):
fixed_calls = []
for call in normalized["tool_calls"]:
call_copy = dict(call)
func = call_copy.get("function")
if isinstance(func, dict):
func_copy = dict(func)
args_val = func_copy.get("arguments")
if not isinstance(args_val, (str, bytes)):
try:
func_copy["arguments"] = json.dumps(args_val, ensure_ascii=False)
except Exception:
func_copy["arguments"] = str(args_val)
call_copy["function"] = func_copy
fixed_calls.append(call_copy)
normalized["tool_calls"] = fixed_calls
return normalized
if not messages:
return AgentLoopOutput(
prompt_ids=[],
response_ids=[],
response_mask=[],
num_turns=0,
metrics={},
)
messages = [_normalize_message(m) for m in messages]
num_turns = 0
for i in range(len(messages)):
if messages[i].get("role") == "system":
continue
# parallel tool calls are in single turn
if i == 0 or messages[i].get("role") != messages[i - 1].get("role"):
num_turns += 1
prompt_ids = []
response_ids = []
response_mask = []
chat_list = []
loop = asyncio.get_running_loop()
# system_prompt_prefix_ids = self.tokenizer.apply_chat_template([{}], add_generation_prompt=False, tokenize=True)
i = 0
try:
while i < len(messages):
if messages[i].get("role") == "system":
chat_list.append(messages[i])
i += 1
continue
# initial chat completion
if messages[i].get("role") == "user":
if i == 0 or messages[i - 1].get("role") == "system":
chat_list.append(messages[i])
prompt_ids = await loop.run_in_executor(
None,
lambda: tokenizer.apply_chat_template(
chat_list,
tools=tools,
add_generation_prompt=True,
tokenize=True,
),
)
else:
chat_list.append(messages[i])
cur_response_ids = await loop.run_in_executor(
None,
lambda: tokenizer.apply_chat_template(
chat_list,
add_generation_prompt=False,
tokenize=True,
),
)
response_ids += cur_response_ids
response_mask += [0] * len(cur_response_ids)
chat_list = []
i += 1
continue
# assistant message
if messages[i].get("role") == "assistant":
chat_list.append(messages[i])
cur_response_ids = await loop.run_in_executor(
None,
lambda: tokenizer.apply_chat_template(
chat_list,
add_generation_prompt=False,
tokenize=True,
),
)
chat_list = []
response_ids += cur_response_ids
response_mask += [1] * len(cur_response_ids)
i += 1
continue
# follow up chat completion with tool response:
if messages[i].get("role") == "tool":
last_assistant_message = messages[i - 1]
chat_list.append(last_assistant_message)
token_assistant = await loop.run_in_executor(
None,
lambda: tokenizer.apply_chat_template(
chat_list,
add_generation_prompt=False,
tokenize=True,
),
)
while i < len(messages) and messages[i].get("role") == "tool":
chat_list.append(messages[i])
i += 1
token_assistant_tool = await loop.run_in_executor(
None,
lambda: tokenizer.apply_chat_template(
chat_list,
add_generation_prompt=False,
tokenize=True,
),
)
tool_response_ids = token_assistant_tool[len(token_assistant):]
chat_list = []
response_ids += tool_response_ids
response_mask += [0] * len(tool_response_ids)
except Exception as e:
raise Exception(f"Failed to convert messages to agentloop_output: {messages}.Exception is: {e}")
max_response_length = min(response_length, len(response_ids))
output = AgentLoopOutput(
prompt_ids=prompt_ids,
response_ids=response_ids[:max_response_length],
response_mask=response_mask[:max_response_length],
num_turns=num_turns,
metrics={},
)
return output
def get_agent_tool_env_and_servers(tool_config: Dict[str, Any] = None) -> tuple[Dict[str, Any], List[str]]:
if not tool_config or not tool_config.get("url") or not tool_config.get("authorization"):
tool_config["url"] = os.getenv("MCP_SERVER_URL")
tool_config["authorization"] = f"Bearer {os.getenv('MCP_SERVER_TOKEN')}"
url = tool_config.get("url")
authorization = tool_config.get("authorization")
mcp_servers_str = tool_config.get("mcp_servers", "")
if not url or not authorization:
raise ValueError("url, Authorization are required. Please set MCP_SERVER_URL and MCP_SERVER_TOKEN environment variable \
or provide them in tool_config parameter.")
server_name = tool_config.get('server_name', 'aworld-mcp')
server_type = tool_config.get('type', 'streamable-http')
timeout = tool_config.get('timeout', 600)
sse_read_timeout = tool_config.get('sse_read_timeout', 600)
client_session_timeout_seconds = tool_config.get('client_session_timeout_seconds', 600)
mcp_config = {
"mcpServers": {
server_name: {
"type": server_type,
"url": url,
"headers": {
"Authorization": authorization,
"MCP_SERVERS": mcp_servers_str,
},
"timeout": timeout,
"sse_read_timeout": sse_read_timeout,
"client_session_timeout_seconds": client_session_timeout_seconds,
}
}
}
servers = list(server_name for server_name in mcp_config.get("mcpServers", {}).keys())
return mcp_config, servers
@@ -0,0 +1,4 @@
from swift.trainers import TrainerFactory
TrainerFactory.TRAINER_MAPPING["aworld_grpo"] = 'train.examples.train_gaia_with_aworld_swift.AworldTrainer'
TrainerFactory.TRAINING_ARGS_MAPPING["aworld_grpo"] = 'train_gaia_with_aworld_swift.trainers.GRPOConfig'
@@ -0,0 +1,19 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
from typing import Union
from aworld.agents.llm_agent import Agent
from aworld.core.agent.swarm import Swarm
from train.adapter.swift.aworld_agent_trainer import AworldTrainer
GAIA_SYSTEM_PROMPT = """
You are an all-capable AI assistant, aimed at solving any task presented by the user.
"""
class GaiaTrainer(AworldTrainer):
def build_agents(self) -> Union[Agent, Swarm]:
return Agent(
name="gaia_super_agent",
system_prompt=GAIA_SYSTEM_PROMPT,
)
@@ -0,0 +1,135 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import re
import string
from typing import List
from swift.plugin import ORM, orms, rm_plugins
from swift.utils import get_logger
logger = get_logger()
"""
Step 1: Define a Reward Class
Implement your custom reward calculation logic within the __call__ method.
The method accepts the model's output completions and dataset columns (passed as kwargs) as input parameters.
Step 2: Register the Reward Class in orms
For example:
python orms['external_math_acc'] = MathAccuracy
Step 3: Configure the Arguments
Use the following arguments when running the script:
bash --plugin /path/to/plugin.py --reward_funcs external_math_acc
"""
class GaiaAnswerMatch(ORM):
def __call__(self, completions, solution, **kwargs) -> List[float]:
pattern = r'<answer>(.*?)</answer>'
rewards = []
logger.info(f"GaiaAnswerMatch|completions:{completions}, comp_match:{solution}")
for content, sol in zip(completions, solution):
comp_match = re.search(pattern, content, re.DOTALL | re.MULTILINE)
logger.info(f"GaiaAnswerMatch|content:{content}, comp_match:{comp_match}, sol:{sol}")
if not comp_match:
rewards.append(0.0)
continue
comp_answer = comp_match.group(1).strip()
if question_scorer(comp_answer, sol):
rewards.append(1.0)
else:
rewards.append(0.0)
return rewards
class GaiaFormat(ORM):
def __call__(self, completions, **kwargs) -> List[float]:
"""Reward function that checks if the completion has a specific format."""
pattern = r'<answer>[\s\S]*?</answer>'
matches = [re.search(pattern, content, re.DOTALL | re.MULTILINE) for content in completions]
reward = [0.1 if match else 0.0 for match in matches]
return reward
orms['external_gaia_answer_reward'] = GaiaAnswerMatch
orms['external_gaia_format_reward'] = GaiaFormat
def split_string(
s: str,
char_list: list[str] = [",", ";"],
) -> list[str]:
pattern = f"[{''.join(char_list)}]"
return re.split(pattern, s)
def normalize_str(input_str, remove_punct=True) -> str:
no_spaces = re.sub(r"\s", "", input_str)
# Remove punctuation, if specified.
if remove_punct:
translator = str.maketrans("", "", string.punctuation)
return no_spaces.lower().translate(translator)
else:
return no_spaces.lower()
def normalize_number_str(number_str: str) -> float:
# we replace these common units and commas to allow
# conversion to float
for char in ["$", "%", ","]:
number_str = number_str.replace(char, "")
try:
return float(number_str)
except ValueError:
# print(f"String {number_str} cannot be normalized to number str.")
return float("inf")
def question_scorer(
model_answer: str,
ground_truth: str,
) -> bool:
def is_float(element: any) -> bool:
try:
float(element)
return True
except ValueError:
return False
if model_answer is None:
model_answer = "None"
# if gt is a number
if is_float(ground_truth):
# print(f"Evaluating {model_answer} as a number.")
normalized_answer = normalize_number_str(model_answer)
return normalized_answer == float(ground_truth)
# if gt is a list
elif any(char in ground_truth for char in [",", ";"]):
# question with the fish: normalization removes punct
gt_elems = split_string(ground_truth)
ma_elems = split_string(model_answer)
# check length is the same
if len(gt_elems) != len(ma_elems):
return False
# compare each element as float or str
comparisons = []
for ma_elem, gt_elem in zip(ma_elems, gt_elems):
if is_float(gt_elem):
normalized_ma_elem = normalize_number_str(ma_elem)
comparisons.append(normalized_ma_elem == float(gt_elem))
else:
# we do not remove punct since comparisons can include punct
comparisons.append(
normalize_str(ma_elem, remove_punct=False)
== normalize_str(gt_elem, remove_punct=False)
)
return all(comparisons)
# if gt is a str
else:
return normalize_str(model_answer) == normalize_str(ground_truth)
@@ -0,0 +1,2 @@
- name: gaia_agent
_target_: train.examples.train_gaia_with_aworld_verl.custom_agent_loop.GaiaAgentLoop
@@ -0,0 +1,34 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
from typing import Union
from aworld.agents.llm_agent import Agent
from aworld.config import AgentConfig
from aworld.core.agent.swarm import Swarm
from train.adapter.verl.aworld_agent_loop import AworldAgentLoop
from train.adapter.verl.common import get_agent_tool_env_and_servers
from env.train_env import TranEnv
GAIA_SYSTEM_PROMPT = """
You are an all-capable AI assistant, aimed at solving any task presented by the user.
"""
class GaiaAgentLoop(AworldAgentLoop):
async def build_agents(self) -> Union[Agent, Swarm]:
gaia_env_config, gaia_env_servers = get_agent_tool_env_and_servers()
return Agent(
conf=AgentConfig(
llm_model_name=await self.get_llm_server_model_name(),
llm_base_url=await self.get_llm_server_address(),
llm_api_key="",
),
name="gaia_super_agent",
system_prompt=GAIA_SYSTEM_PROMPT,
# MCP tool configuration for the agent
mcp_config=gaia_env_config,
mcp_servers=gaia_env_servers,
)
@@ -0,0 +1,67 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import argparse
import json
import os
from pathlib import Path
import pandas as pd
def load_gaia_dataset(path: str, split: str = "validation", total_num_dataset: int = 300):
data_dir = Path(path) / split
split_dataset = []
rl_dataset = {
"prompt": [],
"data_source": [],
"ability": [],
"reward_model": [],
"extra_info": [],
"agent_name": [],
}
cnt = 0
with open(data_dir / "metadata.jsonl", "r", encoding="utf-8") as metaf:
lines = metaf.readlines()
for line in lines:
data = json.loads(line)
if data["task_id"] == "0-0-0-0-0":
continue
if data["file_name"]:
data["file_name"] = data_dir / data["file_name"]
split_dataset.append(data)
rl_dataset["prompt"].append(data["Question"])
rl_dataset["extra_info"].append(
{"task_id": data["task_id"], "split": split, "level": data["Level"], "answer": data["Final answer"]}
)
rl_dataset["agent_name"].append("gaia_agent")
rl_dataset["data_source"].append("gaia")
rl_dataset["ability"].append("agi")
rl_dataset["reward_model"].append({"style": "GAIA", "ground_truth": data['Final answer']})
cnt += 1
if cnt >= total_num_dataset:
break
rl_dataset = pd.DataFrame(data=rl_dataset)
return rl_dataset
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GAIA Dataset Generator")
parser.add_argument("--train_size", type=int, default=300, help="Number of training samples")
parser.add_argument("--test_size", type=int, default=100, help="Number of testing samples")
parser.add_argument("--output_dir", default="gaia_data/", help="Directory to save the dataset")
parser.add_argument("--dataset_path", default="./gaia_dataset", help="GAIA dataset path")
args = parser.parse_args()
gaia_dataset_path = args.dataset_path
train_dataset = load_gaia_dataset(path=gaia_dataset_path, split="validation", total_num_dataset=args.train_size)
test_dataset = load_gaia_dataset(path=gaia_dataset_path, split="test", total_num_dataset=args.test_size)
# Make sure the dataset directory exists
os.makedirs(args.output_dir, exist_ok=True)
# Save the datasets to parquet files
train_dataset.to_parquet(os.path.join(args.output_dir, "train.parquet"))
test_dataset.to_parquet(os.path.join(args.output_dir, "test.parquet"))
@@ -0,0 +1,113 @@
import re
import string
from aworld.logs.util import logger
def normalize_number_str(number_str: str) -> float:
# we replace these common units and commas to allow
# conversion to float
for char in ["$", "%", ","]:
number_str = number_str.replace(char, "")
try:
return float(number_str)
except ValueError:
# print(f"String {number_str} cannot be normalized to number str.")
return float("inf")
def split_string(
s: str,
char_list: list[str] = [",", ";"],
) -> list[str]:
pattern = f"[{''.join(char_list)}]"
return re.split(pattern, s)
def normalize_str(input_str, remove_punct=True) -> str:
"""
Normalize a string by:
- Removing all white spaces
- Optionally removing punctuation (if remove_punct is True)
- Converting to lowercase
Parameters:
- input_str: str, the string to normalize
- remove_punct: bool, whether to remove punctuation (default: True)
Returns:
- str, the normalized string
"""
# Remove all white spaces. Required e.g for seagull vs. sea gull
no_spaces = re.sub(r"\s", "", input_str)
# Remove punctuation, if specified.
if remove_punct:
translator = str.maketrans("", "", string.punctuation)
return no_spaces.lower().translate(translator)
else:
return no_spaces.lower()
def question_scorer(
model_answer: str,
ground_truth: str,
) -> bool:
def is_float(element: any) -> bool:
try:
float(element)
return True
except ValueError:
return False
if model_answer is None:
model_answer = "None"
# if gt is a number
if is_float(ground_truth):
# print(f"Evaluating {model_answer} as a number.")
normalized_answer = normalize_number_str(model_answer)
return normalized_answer == float(ground_truth)
# if gt is a list
elif any(char in ground_truth for char in [",", ";"]):
# print(f"Evaluating {model_answer} as a comma separated list.")
# question with the fish: normalization removes punct
gt_elems = split_string(ground_truth)
ma_elems = split_string(model_answer)
# check length is the same
if len(gt_elems) != len(ma_elems):
# warnings.warn(
# "Answer lists have different lengths, returning False.", UserWarning
# )
return False
# compare each element as float or str
comparisons = []
for ma_elem, gt_elem in zip(ma_elems, gt_elems):
if is_float(gt_elem):
normalized_ma_elem = normalize_number_str(ma_elem)
comparisons.append(normalized_ma_elem == float(gt_elem))
else:
# we do not remove punct since comparisons can include punct
comparisons.append(
normalize_str(ma_elem, remove_punct=False)
== normalize_str(gt_elem, remove_punct=False)
)
return all(comparisons)
# if gt is a str
else:
# print(f"Evaluating {model_answer} as a string.")
return normalize_str(model_answer) == normalize_str(ground_truth)
def gaia_reward_func(data_source, solution_str, ground_truth, extra_info=None):
pattern = r'<answer>(.*?)</answer>'
comp_match = re.search(pattern, solution_str, re.DOTALL | re.MULTILINE)
if not comp_match:
return 0.0
else:
comp_answer = comp_match.group(1).strip()
logger.info(f"comp_answer: {comp_answer}, ground_truth: {ground_truth}")
if question_scorer(comp_answer, ground_truth):
return 1.0
else:
return 0.0
@@ -0,0 +1,139 @@
#!/usr/bin/env bash
set -xeuo pipefail
# ================= cluster topology =================
export GPUS_PER_NODE=${SLURM_GPUS_ON_NODE:-${GPUS_PER_NODE:-1}} # GPUs on this node
NNODES=${SLURM_JOB_NUM_NODES:-${NNODES:-1}}
export NNODES
export RAY_NUM_NODES=$NNODES
echo "Using $NNODES nodes and $GPUS_PER_NODE GPUs per node..."
# ================= data/model/tool =================
HDFS_ROOT=${HDFS_ROOT:-$PWD}
DATA_ROOT=${DATA_ROOT:-$PWD}
# Prefer local model if present, otherwise fall back to HF hub path
model_path=${model_path:-$DATA_ROOT/Qwen/Qwen3-4B}
if [ ! -d "$model_path" ]; then
model_path=Qwen/Qwen3-4B
fi
# Use the default output directory produced by create_dataset.py
train_files=$DATA_ROOT/datasets/train.parquet
test_files=$DATA_ROOT/datasets/test.parquet
# =================== custom ===================
path_to_train="/your/path/to/train"
reward_fn_name=gaia_reward_func
reward_fn_file_path=${path_to_train}/examples/train_gaia_with_aworld_verl/metrics/gaia_reward_function.py
# Agent config
agent_loop_config_path=${path_to_train}/examples/train_gaia_with_aworld_verl/agent.yaml
# set dummy_tool_config_path to enable auto_tool_choice
dummy_tool_config_path=${path_to_train}/examples/verl/configs/dummy_tool_config.yaml
# =================== wandb ===================
project_name=gaia
experiment_name=qwe3
default_local_dir=$DATA_ROOT/checkpoint/$experiment_name
# ================= algorithm =================
adv_estimator=grpo
use_kl_in_reward=false
kl_coef=0.0
use_kl_loss=false
kl_loss_coef=0.0
clip_ratio_low=0.2
clip_ratio_high=0.28
max_turns=8
max_prompt_length=1024
max_response_length=2048
actor_lr=1e-6
train_batch_size=1
ppo_mini_batch_size=1
n_resp_per_prompt=1
n_resp_per_prompt_val=1
# =================== logging ===================
export RAY_LOGGING_LEVEL=DEBUG
export HYDRA_FULL_ERROR=1
# ================= performance =================
export NCCL_IBEXT_DISABLE=1
export NCCL_NVLS_ENABLE=1
export NCCL_IB_HCA=mlx5
export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1
export VLLM_USE_V1=1
export VLLM_ATTENTION_BACKEND=FLASH_ATTN
infer_tp=1 # vLLM tensor parallel size
train_sp=1 # Ulysses sequence parallel size for actor
offload=true
actor_max_token_len_per_gpu=$(( (max_prompt_length + max_response_length) * 4 ))
log_prob_max_token_len_per_gpu=$(( actor_max_token_len_per_gpu * 2 ))
train_files="['$train_files']"
test_files="['$test_files']"
python3 -m verl.trainer.main_ppo \
algorithm.adv_estimator=$adv_estimator \
algorithm.use_kl_in_reward=$use_kl_in_reward \
algorithm.kl_ctrl.kl_coef=$kl_coef \
data.train_files="$train_files" \
data.val_files="$test_files" \
data.return_raw_chat=true \
data.train_batch_size=$train_batch_size \
data.max_prompt_length=$max_prompt_length \
data.max_response_length=$max_response_length \
data.filter_overlong_prompts=true \
data.truncation='error' \
actor_rollout_ref.model.path="$model_path" \
actor_rollout_ref.model.use_remove_padding=true \
actor_rollout_ref.model.enable_gradient_checkpointing=true \
actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \
actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \
actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \
actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \
actor_rollout_ref.actor.clip_ratio_c=10.0 \
actor_rollout_ref.actor.optim.lr=$actor_lr \
actor_rollout_ref.actor.use_dynamic_bsz=true \
actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \
actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu \
actor_rollout_ref.actor.ulysses_sequence_parallel_size=$train_sp \
actor_rollout_ref.actor.fsdp_config.param_offload=$offload \
actor_rollout_ref.actor.fsdp_config.optimizer_offload=$offload \
actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$log_prob_max_token_len_per_gpu \
actor_rollout_ref.rollout.name=vllm \
actor_rollout_ref.rollout.mode=async \
actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \
actor_rollout_ref.rollout.multi_turn.max_user_turns=$max_turns \
actor_rollout_ref.rollout.multi_turn.max_assistant_turns=$max_turns \
actor_rollout_ref.rollout.multi_turn.format=hermes \
actor_rollout_ref.rollout.agent.agent_loop_config_path=$agent_loop_config_path \
actor_rollout_ref.rollout.gpu_memory_utilization=0.75 \
actor_rollout_ref.rollout.n=$n_resp_per_prompt \
actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \
actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \
actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \
actor_rollout_ref.rollout.multi_turn.tool_config_path=$dummy_tool_config_path \
custom_reward_function.path="${reward_fn_file_path}"\
custom_reward_function.name="${reward_fn_name}"\
trainer.logger=console \
trainer.project_name=$project_name \
trainer.experiment_name=$experiment_name \
trainer.n_gpus_per_node="$GPUS_PER_NODE" \
trainer.val_before_train=true \
trainer.log_val_generations=50 \
trainer.nnodes="$NNODES" \
trainer.save_freq=-1 \
trainer.default_local_dir="$default_local_dir" \
trainer.test_freq=5 \
trainer.total_epochs=1 "$@"