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
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:
@@ -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:
|
||||
|
||||

|
||||
|
||||
|
||||
## 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 智能体。
|
||||
|
||||
下图说明了环境和训练集群之间的整体架构和交互:
|
||||
|
||||

|
||||
|
||||
|
||||
## 环境构建
|
||||
|
||||
首先,您需要设置智能体工具将要运行的环境。
|
||||
|
||||
选择一台机器(也可以是训练机)。
|
||||
|
||||
机器规格建议:
|
||||
- 为进行容量规划,为每个并发工作进程分配大约 **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(
|
||||
# 对于要训练的agent,llm_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
|
||||
+4
@@ -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'
|
||||
+19
@@ -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,
|
||||
)
|
||||
+135
@@ -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
|
||||
+34
@@ -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,
|
||||
)
|
||||
+67
@@ -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"))
|
||||
+113
@@ -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 "$@"
|
||||
Reference in New Issue
Block a user