chenyiqun/UnityMAS-O
GitHub: chenyiqun/UnityMAS-O
基于 Verl 改编的 LLM 多智能体强化学习优化框架,将单策略 RL 后训练扩展为可配置的多智能体工作流,支持对整个多智能体系统进行联合训练。
Stars: 54 | Forks: 1
# UnityMAS-O
中文版本:[README.zh.md](README.zh.md)
UnityMAS-O 是一个从 [verl](https://github.com/verl-project/verl) 改编而来的 LLM 多智能体强化学习优化框架。它将传统的单策略 RL 后训练扩展为可配置的多智能体工作流:用户可以定义逻辑智能体、工作流执行图、从智能体到物理 LLM 的映射,以及基于节点、轮次或完整轨迹的奖励分配规则。该框架异步执行工作流,收集结构化轨迹,将奖励分配回相应的智能体,并通过 PPO 风格的训练流水线更新每个物理 LLM。
该仓库仍保留了上游 Verl 训练栈。UnityMAS-O 的特定代码主要位于 `verl/experimental/star_ppo/` 和 `examples/star_ppo/` 中。
## 核心思想
UnityMAS-O 不仅仅训练最终的答案模型。它优化的是整个基于 LLM 的多智能体系统。一个任务实例会被扩展为多步结构化轨迹,例如:
```
QA/search: plan -> search -> retrieve(tool) -> summarize -> update -> answer
code: planner -> coder -> verifier(tool) -> reflector -> planner -> ...
math: solver -> verifier -> refiner -> finalizer
```
该框架明确了四个对象:
- **逻辑智能体 (Logical agents)**:工作流中的角色,例如规划者 (planner)、搜索者 (searcher)、总结者 (summarizer)、程序员 (coder)、反思者 (reflector) 和回答者 (answerer)。
- **智能体-LLM 映射 (Agent-LLM mapping)**:从逻辑智能体到物理模型的映射。智能体可以完全共享一个模型,使用完全独立的模型,或者使用部分共享的模型组。
- **工作流轨迹 (Workflow trace)**:每个样本的结构化执行记录,包括智能体输出、工具结果、状态更新、控制流和调试信息。
- **奖励分配器 (Reward allocator)**:负责将最终指标、局部格式奖励、轮次收益或工具反馈分配回具体智能体调用的组件。
这种设计允许同一个工作流在不同的参数共享方案下进行训练。例如,M-ASK 可以使用四个独立的模型组进行训练,或者将所有角色映射到一个 `shared_agent_llm`;代码工作流可以使用三个独立的模型组分别训练规划者、程序员和反思者,或者切换到共享的 LLM 配置。
## 系统架构
运行时执行遵循 Ray 星型拓扑结构:
- 中央控制器负责调度工作流、调用工具、管理状态转换、汇总奖励并协调训练。
- 每个物理 LLM 都有一个模型本地的工作组,用于 rollout、fat tensor 缓存、ready-batch 构建、优势 (advantage) / logprob / value 计算以及 PPO 更新。
- 控制器仅发送轻量级的动作/输出/元数据记录。大型 tensor 保留在生成它们的工作组内部,从而减少跨节点通信。
- `phi: logical agent -> model_id` 决定了哪个物理模型训练缓冲区接收每个奖励和 rollout 记录。
## 代码结构
```
verl/experimental/star_ppo/
main_ppo.py # UnityMAS-O / STAR PPO entry point
ray_trainer.py # multi-engine Ray trainer, workflow execution, reward commit, PPO update
star_fsdp_workers.py # detached actor / async rollout / critic / reward worker
trajectory_buffer.py # model-local trajectory buffer
types.py # basic types such as engine specs
config/ # Hydra configs
star_ppo_trainer.yaml # shared STAR PPO base config
star_code_iterative_plan_code_reflect_trainer.yaml
star_code_iterative_plan_code_reflect_shared_llm_trainer.yaml
star_iterative_plan_search_summary_update_answer_f1_trainer.yaml
star_iterative_plan_search_summary_update_answer_f1_shared_llm_trainer.yaml
star_math_solver_verifier_refiner_finalizer_*.yaml
star_query_decompose_retrieve*_trainer.yaml
workflows/ # workflow runner plugins
base.py # WorkflowRunner interface
schema.py # WorkflowTrace / WorkflowExecutionRecord / RewardAssignment
mask_iterative_workflow.py # M-ASK iterative search workflow
code_iterative_workflow.py # plan-code-reflect code workflow
math_multi_agent_workflow.py # math multi-agent workflow
graph_workflow.py # graph-style workflow support
reward_allocators/ # reward allocation plugins
base.py
mask_turn_level.py
code_turn_level.py
math_final_answer.py
tools/ # tool interfaces
retriever.py # retrieval API pool
code_verifier.py # local code execution/verifier
math_answer.py
prompt_builders.py
datasets/
code_jsonl_dataset.py # code JSON/JSONL/Parquet adapter
math_jsonl_dataset.py # math JSON/JSONL/Parquet adapter
examples/star_ppo/
common/
run_per_node.sh # start Ray head/worker on each node, launch training on rank 0
run_per_node_background.sh # background launcher, logs go to logs/star_ppo/
run_ip_list.sh # launch by IP list
launch_ip_list_background.sh
launch_kubectl_exec_background.sh
code_iterative_workflow/README.md
mask_iterative_workflow/README.md
math_multi_agent/README.md
```
## 环境设置
从一个干净的 `verl` conda 环境开始。实验使用以下配置运行:
```
cd /path/to/UnityMAS-O
# 创建 Python 3.10 环境。printf 前缀用于自动回答 conda 的交互式提示。
printf 'a\na\nyes\n' | conda create -n verl python=3.10
conda activate verl
# 安装 vLLM / SGLang / Megatron-Core 相关依赖。
bash scripts/install_vllm_sglang_mcore_0.7.sh
# 以 editable 模式安装此 repository,以便代码更改直接生效。
pip install --no-deps -e .
# 固定版本。numpy 2.x 和不同版本的 Transformers/TRL 可能会破坏 Verl/vLLM 的兼容性。
pip install "numpy<2.0"
pip uninstall transformers -y
pip install transformers==4.57 --no-cache-dir
pip uninstall -y trl
pip install "trl==0.26.2"
# 可选:remote debugging。
pip install debugpy==1.8.0
```
该环境主要依赖于 Verl、PyTorch、Ray、vLLM/SGLang、Transformers、Hydra/OmegaConf 和 datasets。如果集群镜像已经包含了部分技术栈,仍需检查 `numpy`、`transformers` 和 `trl` 的版本;许多兼容性问题都源于这些包。
在启动运行之前,清理旧的 Ray 进程和 Python worker 通常会有所帮助:
```
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
```
如果您使用 wandb,请通过环境变量传递凭据。不要将密钥写入脚本或配置文件中:
```
export WANDB_API_KEY=""
export WANDB_ENTITY=""
```
## 运行时私有变量
配置和脚本不应对个人路径、wandb 密钥或内部集群地址进行硬编码。在运行实验之前,请在每个节点的启动环境中设置以下变量:
```
# 个人或集群的 storage root。原有的私有 storage root 已替换为此占位符。
export UNITYMAS_ROOT="/path/to/your/storage/root"
# Ray head node 地址。所有节点必须使用相同的 HEAD_IP;只有 RANK 会改变。
export HEAD_IP=""
# wandb。如果您不使用 wandb,请保持未设置。
export WANDB_API_KEY=""
export WANDB_ENTITY=""
# RAG/search workflows 所需的 Retriever endpoint pool。
export RETRIEVAL_API_URLS_JSON='["http://retriever.example.com:8000/retrieve"]'
# 可选:仅当您的集群需要代理来访问外部网络时才设置。
export PROXY_URL="proxy.example.com:3128"
```
`UNITYMAS_ROOT` 用于构建数据、模型、代码库和安装脚本的默认路径。`HEAD_IP`、`RETRIEVAL_API_URLS_JSON` 和 `PROXY_URL` 是特定于集群的,通常在迁移到新环境时需要重置。请勿将真实值提交到代码库中。
## 多节点启动
通用的启动脚本是 `examples/star_ppo/common/run_per_node_background.sh`。在每个节点上运行一次:
- `HEAD_IP`:rank 0 节点的 IP 地址。所有节点上的此地址必须相同。
- `WORLD_SIZE`:节点总数。
- `RANK`:当前节点的 rank。主节点 (head node) 为 `0`;其他节点为 `1..WORLD_SIZE-1`。
- `CONFIG_NAME`:工作流配置名称。
- 其他环境变量用于指定模型、数据、批大小、rollout 设置、超时和调试开关。
Rank 0 启动 Ray head,等待所有节点加入,然后启动训练。非 0 的 rank 启动 Ray worker 并阻塞。
最小化的主节点命令。每个节点都应首先设置 `UNITYMAS_ROOT`;如果您使用 wandb,请在此处一并设置 `WANDB_API_KEY` 和 `WANDB_ENTITY`:
```
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP=""
export WANDB_API_KEY=""
export WANDB_ENTITY=""
RANK=0 HEAD_IP="${HEAD_IP}" WORLD_SIZE=4 \
CONFIG_NAME=star_iterative_plan_search_summary_update_answer_f1_trainer \
bash examples/star_ppo/common/run_per_node_background.sh
```
Worker 节点:
```
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP=""
export WANDB_API_KEY=""
export WANDB_ENTITY=""
RANK=1 HEAD_IP="${HEAD_IP}" WORLD_SIZE=4 \
CONFIG_NAME=star_iterative_plan_search_summary_update_answer_f1_trainer \
bash examples/star_ppo/common/run_per_node_background.sh
```
日志写入位置:
```
logs/star_ppo/run_rank_.log
```
## 主要工作流
| 工作流 | 配置 | 逻辑智能体 | 典型奖励 |
| --- | --- | --- | --- |
| 反思式代码 (Reflective Code) | `star_code_iterative_plan_code_reflect_trainer` | 规划者、程序员、反思者;验证器作为工具 | 第 0 轮的 verifier pass score,后续轮次的 pass-score delta,加上格式奖励 |
| 共享反思式代码 | `star_code_iterative_plan_code_reflect_shared_llm_trainer` | 规划者/程序员/反思者共享一个物理 LLM | 同上 |
| M-ASK 迭代搜索 | `star_iterative_plan_search_summary_update_answer_f1_trainer` | 规划/回答共享推理 LLM;搜索/总结/更新独立 | 规划/回答使用绝对 F1;搜索/总结/更新使用 F1 delta |
| M-ASK 共享 | `star_iterative_plan_search_summary_update_answer_f1_shared_llm_trainer` | 所有搜索工作流角色共享一个物理 LLM | 同上 |
| 数学多智能体 | `star_math_solver_verifier_refiner_finalizer_trainer` | 求解器、验证器、精炼器、终结器 | 最终答案准确率 + 格式奖励 |
| 查询分解 RAG | `star_query_decompose_retrieve_answer_f1_trainer` 及相关配置 | 查询分解器、回答者、可选的证据/总结者 | 最终答案 F1 + 节点级格式奖励 |
## 任务与基准测试
当前代码库专注于可验证的工作流,其中任务奖励可以通过答案匹配、检索 QA 指标、可执行测试或基于规则的评分器来计算。下表将已完成/已报告的实验与已实现或计划在未来扩展的任务区分开来。
| 状态 | 任务族 | 工作流 / 设置 | 基准测试或数据集 | 主要指标 / 备注 |
| --- | --- | --- | --- | --- |
| 已完成 / 已报告 | 检索增强 QA | 查询分解 -> 检索 -> 回答;查询分解 -> 检索 -> 证据 -> 回答 | Natural Questions (NQ), HotpotQA | 归一化答案 F1;跨模型规模比较多智能体 RL 前后的效果 |
| 已完成 / 已报告 | 迭代智能体搜索 | M-ASK:规划 -> 搜索 -> 总结 -> 更新 -> 回答;独立参数与共享参数变体 | Natural Questions (NQ), HotpotQA | 归一化答案 F1;轮次级 F1 和 F1-delta 奖励;包含 HotpotQA 共享与独立 3B 对比 |
| 已完成 / 已报告 | 反思式代码生成 | 迭代 规划 -> 编写代码 -> 验证(工具) -> 反思,最多三轮 | DeepCoder 风格的可验证编程任务;基于 TACO-Verified、PrimeIntellect SYNTHETIC-1 和 LiveCodeBench v5 处理的训练/测试集划分 | 可执行测试通过率、留出集全部通过率以及平均验证器轮数 |
| 已实现 / 扩展中 | 数学推理 | 求解器 -> 验证器 -> 精炼器 -> 终结器 | DeepScaleR 风格的训练数据;MATH-500, AIME24, AIME25, AIME26, AMC23 | 最终答案准确率和格式奖励;存在配置/数据集适配器,正在进行更广泛的实验扩展 |
| 计划中 / 进行中 | 具身智能体任务 | 工具/环境交互工作流 | ALFWorld | 任务成功率/完成率;旨在测试长周期环境交互 |
| 计划中 / 进行中 | 网络交互任务 | 搜索/浏览/购物交互工作流 | WebShop | WebShop 奖励或成功分数;旨在测试网络动作和工具使用的信用分配 |
| 计划中 / 进行中 | 软件工程任务 | 多智能体问题定位、编码、验证和修复 | SWE-bench | 已解决 Issue 率/测试通过率;旨在测试长周期的代码库级工作流 |
## 示例 1:反思式代码生成
此配置训练三个非共享的 LLM 智能体:
- `planner_agent` -> `planner_llm`
- `coder_agent` -> `coder_llm`
- `reflection_agent` -> `reflection_llm`
数据集可以是 JSON、JSONL 或 Parquet。常见字段有:
```
{
"uid": "example/0",
"source": "codeforces",
"problem": "problem statement ...",
"starter_code": "",
"tests": "[{\"input\":\"1\\n2 1\",\"output\":\"1 2\"}]"
}
```
三节点示例:在三个节点上运行相同的命令,只需将 `RANK` 分别更改为 `0`、`1` 和 `2`。
```
cd /path/to/UnityMAS-O
conda activate verl
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP=""
export WANDB_API_KEY=""
export WANDB_ENTITY=""
RANK=0 HEAD_IP="${HEAD_IP}" WORLD_SIZE=3 \
CONFIG_NAME=star_code_iterative_plan_code_reflect_trainer \
PROJECT_NAME="STAR-Code" \
EXPERIMENT_NAME="deepcoder_marti_iterative_plan_code_reflect_3xQwen3_4B_no_think_sp4" \
TRAIN_JSONL="${UNITYMAS_ROOT}/datasets/code_datasets/DeepCoder-Preview-Dataset/processed_marti_jsonl/train_shuffled.jsonl" \
VAL_JSONL="${UNITYMAS_ROOT}/datasets/code_datasets/DeepCoder-Preview-Dataset/processed_marti_jsonl/test_shuffled.jsonl" \
AGENT_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen3-4B" \
PLANNER_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen3-4B" \
CODER_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen3-4B" \
REFLECTION_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen3-4B" \
ACTOR_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen3-4B" \
QWEN_ENABLE_THINKING=false \
GEN_BATCH_SIZE=64 \
VAL_BATCH_SIZE=64 \
VAL_MAX_BATCHES=-1 \
VAL_BEFORE_TRAIN=true \
STAR_MAX_INFLIGHT_QUERIES=64 \
STAR_MAX_PARALLEL_ROLLOUTS_PER_MODEL=64 \
STAR_LLM_MICROBATCH_MAX_SIZE=64 \
STAR_LLM_MICROBATCH_MAX_WAIT_MS=1000 \
ACTOR_PPO_MINI_BATCH_SIZE=64 \
ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=1 \
CRITIC_PPO_MICRO_BATCH_SIZE_PER_GPU=1 \
ROLLOUT_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=1 \
REF_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=1 \
ROLLOUT_TENSOR_MODEL_PARALLEL_SIZE=4 \
ROLLOUT_GPU_MEMORY_UTILIZATION=0.40 \
ROLLOUT_PROMPT_LENGTH=8192 \
ROLLOUT_RESPONSE_LENGTH=2048 \
ROLLOUT_MAX_MODEL_LEN=10240 \
ROLLOUT_MAX_NUM_SEQS=64 \
DATA_MAX_PROMPT_LENGTH=8192 \
STAR_PER_INFER_PROMPT_MAX_TOKENS=7680 \
CODE_MAX_TURNS=3 \
CODE_STOP_ON_ALL_PASSED=true \
CODE_VERIFY_TIMEOUT_SECONDS=1.0 \
CODE_VERIFY_DEFAULT_CHECKER_TYPE=auto \
CODE_VERIFY_MAX_TESTS_PER_EXAMPLE=8 \
CODE_VERIFIER_FAIL_OPEN=false \
STAR_QUERY_TIMEOUT_SECONDS=420 \
STAR_WORKFLOW_BATCH_TIMEOUT_SECONDS=900 \
STAR_RAY_GET_TIMEOUT_SECONDS=600 \
STAR_WORKER_CALL_TIMEOUT_SECONDS=600 \
STAR_LLM_TIMEOUT_SECONDS=900 \
STAR_VAL_PROGRESS_EVERY=1 \
STAR_WORKFLOW_DEBUG=true \
STAR_WORKFLOW_DEBUG_EVERY_N_BATCHES=10 \
STAR_WORKFLOW_DEBUG_SAMPLE_INDEX=0 \
STAR_WORKFLOW_DEBUG_MAX_CHARS=4000 \
STAR_VAL_DEBUG=true \
STAR_VAL_DEBUG_MAX_CHARS=4000 \
STAR_TOOL_TIMEOUT_SECONDS=0 \
bash examples/star_ppo/common/run_per_node_background.sh \
actor_rollout_ref.model.use_remove_padding=true \
critic.model.use_remove_padding=true \
actor_rollout_ref.actor.ulysses_sequence_parallel_size=4 \
actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=4 \
actor_rollout_ref.ref.ulysses_sequence_parallel_size=4 \
actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=4 \
critic.ulysses_sequence_parallel_size=4 \
critic.model.fsdp_config.ulysses_sequence_parallel_size=4
```
重要开关:
- `CODE_MAX_TURNS`:规划-编码-验证-反思的最大轮数。
- `CODE_STOP_ON_ALL_PASSED`:一旦所有验证器测试通过,则提前停止。
- `CODE_VERIFY_TIMEOUT_SECONDS`:一次代码执行的超时时间。
- `CODE_VERIFIER_FAIL_OPEN=false`:验证器失败是否应视为通过。对于代码训练,通常首选 `false`。
- `STAR_PER_INFER_PROMPT_MAX_TOKENS`:每个智能体 prompt 的截断限制。
## 示例 2:M-ASK 迭代搜索,四个模型组
此配置包含五个逻辑智能体和四个物理 LLM:
- `planning_agent` 和 `answer_agent` 共享 `reasoning_agent_llm`
- `search_agent` 使用独立的 LLM
- `summary_agent` 使用独立的 LLM
- `update_agent` 使用独立的 LLM
默认情况下,数据读取自:
```
DATASET_ROOT//train_verl.parquet
DATASET_ROOT//test_verl.parquet
```
您可以使用 `TRAIN_PARQUET` 和 `VAL_PARQUET` 覆盖此项。
四节点示例:在四个节点上运行相同的命令,只需将 `RANK` 分别更改为 `0`、`1`、`2` 和 `3`。
```
cd /path/to/UnityMAS-O
conda activate verl
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP=""
export WANDB_API_KEY=""
export WANDB_ENTITY=""
RANK=0 HEAD_IP="${HEAD_IP}" WORLD_SIZE=4 \
CONFIG_NAME=star_iterative_plan_search_summary_update_answer_f1_trainer \
DATASET_NAME="hotpotqa" \
STAR_RETRIEVER_RANDOM_ENDPOINT=true \
RETRIEVAL_API_URLS_JSON='["http://retriever-0.example.com:8000/retrieve","http://retriever-1.example.com:8000/retrieve"]' \
PROJECT_NAME="M-ASK" \
EXPERIMENT_NAME="hotpotqa_M-ASK_f1_4x7B" \
REASONING_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen2.5-7B-Instruct" \
SEARCH_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen2.5-7B-Instruct" \
SUMMARY_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen2.5-7B-Instruct" \
UPDATE_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen2.5-7B-Instruct" \
GEN_BATCH_SIZE=128 \
STAR_MAX_INFLIGHT_QUERIES=128 \
STAR_MAX_PARALLEL_ROLLOUTS_PER_MODEL=32 \
ACTOR_PPO_MINI_BATCH_SIZE=128 \
ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=1 \
CRITIC_PPO_MICRO_BATCH_SIZE_PER_GPU=1 \
ROLLOUT_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=1 \
REF_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=1 \
ROLLOUT_TENSOR_MODEL_PARALLEL_SIZE=1 \
ROLLOUT_GPU_MEMORY_UTILIZATION=0.20 \
MASK_MAX_TURNS=3 \
MASK_STOP_ON_SEARCH_END=true \
STAR_QUERY_TIMEOUT_SECONDS=600 \
STAR_WORKFLOW_BATCH_TIMEOUT_SECONDS=900 \
STAR_RAY_GET_TIMEOUT_SECONDS=300 \
STAR_WORKER_CALL_TIMEOUT_SECONDS=300 \
STAR_LLM_TIMEOUT_SECONDS=300 \
STAR_VAL_PROGRESS_EVERY=1 \
STAR_WORKFLOW_DEBUG=true \
STAR_WORKFLOW_DEBUG_EVERY_N_BATCHES=10 \
STAR_WORKFLOW_DEBUG_SAMPLE_INDEX=0 \
STAR_WORKFLOW_DEBUG_MAX_CHARS=160 \
bash examples/star_ppo/common/run_per_node_background.sh
```
M-ASK 奖励分配:
- 规划智能体:初始答案 `a0` 的绝对 F1。
- 回答智能体:每个临时答案 `at` 的绝对 F1。
- 搜索/总结/更新:共享边际改进 `F1(at) - F1(at-1)`。
- 当搜索输出 `` 时,该搜索步骤的任务奖励为 0。
## 示例 3:M-ASK 共享 LLM,一个模型组
此配置将规划/搜索/总结/更新/回答映射到单个 `shared_agent_llm`。这对于参数共享研究、节省资源的运行以及快速的小模型实验非常有用。
```
cd /path/to/UnityMAS-O
conda activate verl
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP=""
export WANDB_API_KEY=""
export WANDB_ENTITY=""
RANK=0 HEAD_IP="${HEAD_IP}" WORLD_SIZE=1 \
CONFIG_NAME=star_iterative_plan_search_summary_update_answer_f1_shared_llm_trainer \
DATASET_NAME="hotpotqa" \
STAR_RETRIEVER_RANDOM_ENDPOINT=true \
RETRIEVAL_API_URLS_JSON='["http://retriever-0.example.com:8000/retrieve","http://retriever-1.example.com:8000/retrieve"]' \
PROJECT_NAME="M-ASK" \
EXPERIMENT_NAME="hotpotqa_M-ASK_f1_3B_shared" \
SHARED_MODEL_PATH="${UNITYMAS_ROOT}/base_models/Qwen/Qwen2.5-3B-Instruct" \
GEN_BATCH_SIZE=128 \
STAR_MAX_INFLIGHT_QUERIES=128 \
STAR_MAX_PARALLEL_ROLLOUTS_PER_MODEL=32 \
ACTOR_PPO_MINI_BATCH_SIZE=48 \
ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=6 \
CRITIC_PPO_MICRO_BATCH_SIZE_PER_GPU=6 \
ROLLOUT_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=6 \
REF_LOGPROB_MICRO_BATCH_SIZE_PER_GPU=6 \
ROLLOUT_TENSOR_MODEL_PARALLEL_SIZE=1 \
ROLLOUT_GPU_MEMORY_UTILIZATION=0.20 \
MASK_MAX_TURNS=3 \
MASK_STOP_ON_SEARCH_END=true \
STAR_QUERY_TIMEOUT_SECONDS=600 \
STAR_WORKFLOW_BATCH_TIMEOUT_SECONDS=900 \
STAR_RAY_GET_TIMEOUT_SECONDS=300 \
STAR_WORKER_CALL_TIMEOUT_SECONDS=300 \
STAR_LLM_TIMEOUT_SECONDS=300 \
STAR_VAL_PROGRESS_EVERY=1 \
STAR_WORKFLOW_DEBUG=true \
STAR_WORKFLOW_DEBUG_EVERY_N_BATCHES=10 \
STAR_WORKFLOW_DEBUG_SAMPLE_INDEX=0 \
STAR_WORKFLOW_DEBUG_MAX_CHARS=160 \
bash examples/star_ppo/common/run_per_node_background.sh
```
## 常用环境变量
| 变量 | 用途 |
| --- | --- |
| `CONFIG_NAME` | Hydra 配置名称,不带 `.yaml` |
| `PROJECT_NAME`, `EXPERIMENT_NAME` | wandb/控制台跟踪名称 |
| `RANK`, `HEAD_IP`, `WORLD_SIZE` | 多节点 Ray 启动参数 |
| `GPUS_PER_NODE`, `CPUS_PER_NODE` | 每节点资源声明 |
| `AGENT_MODEL_PATH` | 适用于大多数配置的通用回退模型路径 |
| `ACTOR_MODEL_PATH`, `ACTOR_TOKENIZER_PATH` | Verl actor/ref/critic 的基础模型和 tokenizer |
| `PLANNER_MODEL_PATH`, `CODER_MODEL_PATH`, `REFLECTION_MODEL_PATH` | 代码工作流的三个模型路径 |
| `REASONING_MODEL_PATH`, `SEARCH_MODEL_PATH`, `SUMMARY_MODEL_PATH`, `UPDATE_MODEL_PATH` | 非共享 M-ASK 配置的模型路径 |
| `SHARED_MODEL_PATH` | 共享 LLM 配置的模型路径 |
| `TRAIN_JSONL`, `VAL_JSONL` | 代码/数学 JSONL 数据路径 |
| `TRAIN_PARQUET`, `VAL_PARQUET`, `DATASET_ROOT`, `DATASET_NAME` | QA/搜索 Verl 格式 parquet 数据路径 |
| `GEN_BATCH_SIZE`, `VAL_BATCH_SIZE` | rollout 生成批大小和验证批大小 |
| `ACT_PPO_MINI_BATCH_SIZE` | PPO mini-batch 大小 |
| `ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU` | actor micro-batch 大小 |
| `ROLLOUT_TENSOR_MODEL_PARALLEL_SIZE` | vLLM tensor parallel 大小 |
| `ROLLOUT_GPU_MEMORY_UTILIZATION` | vLLM GPU 显存占用比例 |
| `ROLLOUT_PROMPT_LENGTH`, `ROLLOUT_RESPONSE_LENGTH`, `ROLLOUT_MAX_MODEL_LEN` | rollout 长度控制 |
| `STAR_MAX_INFLIGHT_QUERIES` | 控制器执行的并发查询数 |
| `STAR_MAX_PARALLEL_ROLLOUTS_PER_MODEL` | 每个 `model_id` 的并发 rollout 数 |
| `STAR_LLM_MICROBATCH_MAX_SIZE`, `STAR_LLM_MICROBATCH_MAX_WAIT_MS` | LLM 请求微批处理 |
| `STAR_QUERY_TIMEOUT_SECONDS` | 一次查询工作流的超时时间 |
| `STAR_WORKFLOW_BATCH_TIMEOUT_SECONDS` | 一个工作流批次的超时时间 |
| `STAR_RAY_GET_TIMEOUT_SECONDS`, `STAR_WORKER_CALL_TIMEOUT_SECONDS` | Ray/worker 调用超时 |
| `STAR_LLM_TIMEOUT_SECONDS`, `STAR_TOOL_TIMEOUT_SECONDS` | LLM/工具调用超时 |
| `STAR_WORKFLOW_DEBUG`, `STAR_VAL_DEBUG` | 打印工作流轨迹调试信息 |
## 数据格式
### 代码 JSONL
`CodeJsonlDataset` 读取 `problem/question/query` 作为问题描述,`tests/test_cases/answer/label/reward_model/extra_info.*` 作为测试用例,读取 `starter_code/extra_info.starter_code` 作为起始代码。
最小化示例:
```
{"uid":"code/0","problem":"Write a function ...","starter_code":"","tests":[{"input":"1\n","output":"1\n"}]}
```
`tests` 也可以存储为 JSON 字符串。
### QA / 搜索 Parquet
QA/搜索配置默认使用 Verl 格式的 parquet。常见字段包括:
- `question` / `query` / `problem` / `extra_info.question`
- `answer` / `ground_truth` / `extra_info.answer` / `reward_model.ground_truth`
检索工具通过 `RETRIEVAL_API_URLS_JSON` 接收一个或多个 HTTP endpoint。每个 endpoint 应公开一个 `/retrieve` API,并返回供工作流运行器使用的候选文档。
### 数学 JSONL
`MathJsonlDataset` 支持 JSON、JSONL 和 Parquet。它读取 `question/problem/query` 作为问题,读取 `answer/ground_truth/target/reward_model.ground_truth/solution` 作为答案,并自动推断 `data_source` 以进行特定数据集的验证指标计算。
## 添加工作流
添加新任务通常需要三种类型的文件:
1. 在 `verl/experimental/star_ppo/workflows/` 下实现一个 `WorkflowRunner`。
2. 在 `verl/experimental/star_ppo/reward_allocators/` 下实现一个 `RewardAllocator`。
3. 在 `verl/experimental/star_ppo/config/` 下添加一个 Hydra YAML 配置,声明 `trainer.llm_engines`、智能体到 `model_id` 的映射、运行器、奖励分配器、工具和数据路径。
核心 `WorkflowRunner` 接口:
```
class WorkflowRunner:
async def run_batch(self, batch: DataProto, epoch: int) -> tuple[DataProto, dict[str, float]]:
...
```
核心 `RewardAllocator` 接口:
```
class RewardAllocator:
def allocate(self, trace: WorkflowTrace) -> tuple[list[RewardAssignment], dict[str, float]]:
...
```
关键约定:
- 每个可训练的 LLM 调用都应创建一个 `WorkflowExecutionRecord`,并保留其对应的 thin/fat 轨迹 ID。
- 工具节点可以出现在轨迹中,但不需要进入 PPO 训练缓冲区。
- 奖励分配器最终将标量奖励绑定到具体的 `WorkflowExecutionRecord`。
- 只要奖励可以通过 `traj_id` 提交回正确的缓冲区,PPO 训练器就不需要理解特定于工作流的语义。
## 调试
检查后台日志:
```
tail -f logs/star_ppo/run_rank0_*.log
```
检查 Ray 集群:
```
ray status
```
常见问题:
- **非主节点未加入**:确保所有节点使用相同的 `HEAD_IP`、唯一的 `RANK`、正确的 `WORLD_SIZE` 以及可访问的端口 `6379/8265`。
- **训练在等待活动节点时挂起**:至少有一个 worker 未能启动;请先检查相应 rank 的日志。
- **vLLM OOM**:降低 `ROLLOUT_GPU_MEMORY_UTILIZATION`、`ROLLOUT_MAX_NUM_SEQS`、`ROLLOUT_MAX_NUM_BATCHED_TOKENS` 或 `STAR_MAX_PARALLEL_ROLLOUTS_PER_MODEL`。
- **Prompt 过长**:降低 `STAR_PER_INFER_PROMPT_MAX_TOKENS` 和 `DATA_MAX_PROMPT_LENGTH`,或启用/调整工作流状态截断。
- **验证器太慢**:减少 `CODE_VERIFY_MAX_TESTS_PER_EXAMPLE`,增加 `CODE_VERIFY_TIMEOUT_SECONDS`,并检查测试用例的大小限制。
- **检索不稳定**:添加更多 `RETRIEVAL_API_URLS_JSON` endpoint,设置 `STAR_RETRIEVER_RANDOM_ENDPOINT=true`,并检查检索服务器的超时设置。
- **调试输出过多**:禁用 `STAR_WORKFLOW_DEBUG` / `STAR_VAL_DEBUG`,或降低 `STAR_WORKFLOW_DEBUG_MAX_CHARS`。
## 与 Verl 的关系
UnityMAS-O 复用了 Verl 的核心训练基础设施,包括 Ray 分布式执行、FSDP/FSDP2 worker、actor/ref/critic、vLLM rollout、PPO 更新、跟踪和 checkpointing。在此基础之上,本代码库添加了用于多智能体工作流的控制器、路由、轨迹、奖励分配、模型本地轨迹缓冲区和多 LLM 引擎配置。
如果您只需要原始 Verl 的单策略 PPO/GRPO/SFT 功能,您仍然可以使用 `verl/trainer/` 和 `examples/ppo_trainer/` 下的原始入口点。要训练多智能体工作流,请使用 `verl.experimental.star_ppo.main_ppo` 和 `examples/star_ppo/` 下的脚本。
## 技术报告与结果
相应的技术报告标题为:
```
UnityMAS-O: A General RL Optimization Framework for LLM-Based Multi-Agent Systems
```
该报告评估了三个工作流族:QA/搜索、M-ASK 迭代搜索和反思式代码生成。主要发现是 UnityMAS-O 能够将这些手动设计的多智能体工作流转化为可训练的 MARL 问题。训练后,QA F1、代码全通过率和代码验证轮次均得到明显改善,同时该框架还支持在共享参数和独立模型设置之间进行受控比较。
在 QA 任务中,经过训练的工作流在不同的模型规模和检索流水线上都有持续的改进:
共享参数的 M-ASK 变体收敛速度稍慢,但在 HotpotQA 上可以接近独立的四模型组设置:
对于代码生成,规划-编码-验证-反思工作流在训练后显著提高了留出集的全通过率:
同样的实验还表明,经过训练的代码工作流能更早地通过验证器,从而减少了平均验证轮次:
## 核心思想
UnityMAS-O 不仅仅训练最终的答案模型。它优化的是整个基于 LLM 的多智能体系统。一个任务实例会被扩展为多步结构化轨迹,例如:
```
QA/search: plan -> search -> retrieve(tool) -> summarize -> update -> answer
code: planner -> coder -> verifier(tool) -> reflector -> planner -> ...
math: solver -> verifier -> refiner -> finalizer
```
该框架明确了四个对象:
- **逻辑智能体 (Logical agents)**:工作流中的角色,例如规划者 (planner)、搜索者 (searcher)、总结者 (summarizer)、程序员 (coder)、反思者 (reflector) 和回答者 (answerer)。
- **智能体-LLM 映射 (Agent-LLM mapping)**:从逻辑智能体到物理模型的映射。智能体可以完全共享一个模型,使用完全独立的模型,或者使用部分共享的模型组。
- **工作流轨迹 (Workflow trace)**:每个样本的结构化执行记录,包括智能体输出、工具结果、状态更新、控制流和调试信息。
- **奖励分配器 (Reward allocator)**:负责将最终指标、局部格式奖励、轮次收益或工具反馈分配回具体智能体调用的组件。
这种设计允许同一个工作流在不同的参数共享方案下进行训练。例如,M-ASK 可以使用四个独立的模型组进行训练,或者将所有角色映射到一个 `shared_agent_llm`;代码工作流可以使用三个独立的模型组分别训练规划者、程序员和反思者,或者切换到共享的 LLM 配置。
## 系统架构
运行时执行遵循 Ray 星型拓扑结构:
- 中央控制器负责调度工作流、调用工具、管理状态转换、汇总奖励并协调训练。
- 每个物理 LLM 都有一个模型本地的工作组,用于 rollout、fat tensor 缓存、ready-batch 构建、优势 (advantage) / logprob / value 计算以及 PPO 更新。
- 控制器仅发送轻量级的动作/输出/元数据记录。大型 tensor 保留在生成它们的工作组内部,从而减少跨节点通信。
- `phi: logical agent -> model_id` 决定了哪个物理模型训练缓冲区接收每个奖励和 rollout 记录。
## 代码结构
```
verl/experimental/star_ppo/
main_ppo.py # UnityMAS-O / STAR PPO entry point
ray_trainer.py # multi-engine Ray trainer, workflow execution, reward commit, PPO update
star_fsdp_workers.py # detached actor / async rollout / critic / reward worker
trajectory_buffer.py # model-local trajectory buffer
types.py # basic types such as engine specs
config/ # Hydra configs
star_ppo_trainer.yaml # shared STAR PPO base config
star_code_iterative_plan_code_reflect_trainer.yaml
star_code_iterative_plan_code_reflect_shared_llm_trainer.yaml
star_iterative_plan_search_summary_update_answer_f1_trainer.yaml
star_iterative_plan_search_summary_update_answer_f1_shared_llm_trainer.yaml
star_math_solver_verifier_refiner_finalizer_*.yaml
star_query_decompose_retrieve*_trainer.yaml
workflows/ # workflow runner plugins
base.py # WorkflowRunner interface
schema.py # WorkflowTrace / WorkflowExecutionRecord / RewardAssignment
mask_iterative_workflow.py # M-ASK iterative search workflow
code_iterative_workflow.py # plan-code-reflect code workflow
math_multi_agent_workflow.py # math multi-agent workflow
graph_workflow.py # graph-style workflow support
reward_allocators/ # reward allocation plugins
base.py
mask_turn_level.py
code_turn_level.py
math_final_answer.py
tools/ # tool interfaces
retriever.py # retrieval API pool
code_verifier.py # local code execution/verifier
math_answer.py
prompt_builders.py
datasets/
code_jsonl_dataset.py # code JSON/JSONL/Parquet adapter
math_jsonl_dataset.py # math JSON/JSONL/Parquet adapter
examples/star_ppo/
common/
run_per_node.sh # start Ray head/worker on each node, launch training on rank 0
run_per_node_background.sh # background launcher, logs go to logs/star_ppo/
run_ip_list.sh # launch by IP list
launch_ip_list_background.sh
launch_kubectl_exec_background.sh
code_iterative_workflow/README.md
mask_iterative_workflow/README.md
math_multi_agent/README.md
```
## 环境设置
从一个干净的 `verl` conda 环境开始。实验使用以下配置运行:
```
cd /path/to/UnityMAS-O
# 创建 Python 3.10 环境。printf 前缀用于自动回答 conda 的交互式提示。
printf 'a\na\nyes\n' | conda create -n verl python=3.10
conda activate verl
# 安装 vLLM / SGLang / Megatron-Core 相关依赖。
bash scripts/install_vllm_sglang_mcore_0.7.sh
# 以 editable 模式安装此 repository,以便代码更改直接生效。
pip install --no-deps -e .
# 固定版本。numpy 2.x 和不同版本的 Transformers/TRL 可能会破坏 Verl/vLLM 的兼容性。
pip install "numpy<2.0"
pip uninstall transformers -y
pip install transformers==4.57 --no-cache-dir
pip uninstall -y trl
pip install "trl==0.26.2"
# 可选:remote debugging。
pip install debugpy==1.8.0
```
该环境主要依赖于 Verl、PyTorch、Ray、vLLM/SGLang、Transformers、Hydra/OmegaConf 和 datasets。如果集群镜像已经包含了部分技术栈,仍需检查 `numpy`、`transformers` 和 `trl` 的版本;许多兼容性问题都源于这些包。
在启动运行之前,清理旧的 Ray 进程和 Python worker 通常会有所帮助:
```
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
```
如果您使用 wandb,请通过环境变量传递凭据。不要将密钥写入脚本或配置文件中:
```
export WANDB_API_KEY="
| 工作流 | 配置 | 逻辑智能体 | 典型奖励 |
| --- | --- | --- | --- |
| 反思式代码 (Reflective Code) | `star_code_iterative_plan_code_reflect_trainer` | 规划者、程序员、反思者;验证器作为工具 | 第 0 轮的 verifier pass score,后续轮次的 pass-score delta,加上格式奖励 |
| 共享反思式代码 | `star_code_iterative_plan_code_reflect_shared_llm_trainer` | 规划者/程序员/反思者共享一个物理 LLM | 同上 |
| M-ASK 迭代搜索 | `star_iterative_plan_search_summary_update_answer_f1_trainer` | 规划/回答共享推理 LLM;搜索/总结/更新独立 | 规划/回答使用绝对 F1;搜索/总结/更新使用 F1 delta |
| M-ASK 共享 | `star_iterative_plan_search_summary_update_answer_f1_shared_llm_trainer` | 所有搜索工作流角色共享一个物理 LLM | 同上 |
| 数学多智能体 | `star_math_solver_verifier_refiner_finalizer_trainer` | 求解器、验证器、精炼器、终结器 | 最终答案准确率 + 格式奖励 |
| 查询分解 RAG | `star_query_decompose_retrieve_answer_f1_trainer` 及相关配置 | 查询分解器、回答者、可选的证据/总结者 | 最终答案 F1 + 节点级格式奖励 |
## 任务与基准测试
当前代码库专注于可验证的工作流,其中任务奖励可以通过答案匹配、检索 QA 指标、可执行测试或基于规则的评分器来计算。下表将已完成/已报告的实验与已实现或计划在未来扩展的任务区分开来。
| 状态 | 任务族 | 工作流 / 设置 | 基准测试或数据集 | 主要指标 / 备注 |
| --- | --- | --- | --- | --- |
| 已完成 / 已报告 | 检索增强 QA | 查询分解 -> 检索 -> 回答;查询分解 -> 检索 -> 证据 -> 回答 | Natural Questions (NQ), HotpotQA | 归一化答案 F1;跨模型规模比较多智能体 RL 前后的效果 |
| 已完成 / 已报告 | 迭代智能体搜索 | M-ASK:规划 -> 搜索 -> 总结 -> 更新 -> 回答;独立参数与共享参数变体 | Natural Questions (NQ), HotpotQA | 归一化答案 F1;轮次级 F1 和 F1-delta 奖励;包含 HotpotQA 共享与独立 3B 对比 |
| 已完成 / 已报告 | 反思式代码生成 | 迭代 规划 -> 编写代码 -> 验证(工具) -> 反思,最多三轮 | DeepCoder 风格的可验证编程任务;基于 TACO-Verified、PrimeIntellect SYNTHETIC-1 和 LiveCodeBench v5 处理的训练/测试集划分 | 可执行测试通过率、留出集全部通过率以及平均验证器轮数 |
| 已实现 / 扩展中 | 数学推理 | 求解器 -> 验证器 -> 精炼器 -> 终结器 | DeepScaleR 风格的训练数据;MATH-500, AIME24, AIME25, AIME26, AMC23 | 最终答案准确率和格式奖励;存在配置/数据集适配器,正在进行更广泛的实验扩展 |
| 计划中 / 进行中 | 具身智能体任务 | 工具/环境交互工作流 | ALFWorld | 任务成功率/完成率;旨在测试长周期环境交互 |
| 计划中 / 进行中 | 网络交互任务 | 搜索/浏览/购物交互工作流 | WebShop | WebShop 奖励或成功分数;旨在测试网络动作和工具使用的信用分配 |
| 计划中 / 进行中 | 软件工程任务 | 多智能体问题定位、编码、验证和修复 | SWE-bench | 已解决 Issue 率/测试通过率;旨在测试长周期的代码库级工作流 |
## 示例 1:反思式代码生成
此配置训练三个非共享的 LLM 智能体:
- `planner_agent` -> `planner_llm`
- `coder_agent` -> `coder_llm`
- `reflection_agent` -> `reflection_llm`
数据集可以是 JSON、JSONL 或 Parquet。常见字段有:
```
{
"uid": "example/0",
"source": "codeforces",
"problem": "problem statement ...",
"starter_code": "",
"tests": "[{\"input\":\"1\\n2 1\",\"output\":\"1 2\"}]"
}
```
三节点示例:在三个节点上运行相同的命令,只需将 `RANK` 分别更改为 `0`、`1` 和 `2`。
```
cd /path/to/UnityMAS-O
conda activate verl
ray stop --force >/dev/null 2>&1 || true
pkill -9 -f "/miniconda3/envs/verl/bin/python3.10" || true
export UNITYMAS_ROOT="/path/to/your/storage/root"
export HEAD_IP="
共享参数的 M-ASK 变体收敛速度稍慢,但在 HotpotQA 上可以接近独立的四模型组设置:
对于代码生成,规划-编码-验证-反思工作流在训练后显著提高了留出集的全通过率:
同样的实验还表明,经过训练的代码工作流能更早地通过验证器,从而减少了平均验证轮次:
标签:DLL 劫持, PPO, RLHF, 凭据扫描, 多智能体, 大语言模型, 强化学习, 模型优化, 系统调用监控, 逆向工具