从零开始学 verl 框架

"""
- trainer
  - config
    - ppo_trainer.yaml
  - main_ppo.py
"""
import hydra
from omegaconf import DictConfig

@hydra.main(config_path="config", config_name="ppo_trainer", version_base=None)
def main(config):
	# ...
    run_ppo(config)

@ray.remote
class Accumulator:
    def __init__(self):
        self.value = 0

    def add(self, x):
        self.value += x

    def get_value(self):
        return self.value

accumulator = Accumulator.remote()
value_ref = accumulator.get_value.remote()
# ObjectRef(16310a0f0a45af5c9e4edcae28cff4ef250feac40100000001000000)
value = ray.get(value_ref)  
# 0

@ray.remote
class GPUAccumulator(Worker):
    def __init__(self):
        super().__init__()
        self.value = torch.zeros(1, device="cuda") + self.rank

    @register(Dispatch.ONE_TO_ALL)
    def add(self, x):
        self.value += x
        return self.value.cpu()

resource_pool = RayResourcePool([1], use_gpu=True)
class_with_args = RayClassWithInitArgs(cls=GPUAccumulator)
worker_group = RayWorkerGroup(resource_pool, class_with_args)
print(worker_group.add(x=10))

class DataProto:
    batch: TensorDict          # tensor 类型的字段（input_ids, attention_mask 等）
    non_tensor_batch: dict     # 非 tensor 字段（data_source, ground_truth 等）
    meta_info: dict            # 元信息（seqlen 等）

def train(self) -> None:
    """
    标准 RLHF PPO pipeline
    1. 用 actor 根据 prompt 生成回答
    2. 用 actor/ref 计算每个 token 的 logprobs
    3. 用 reward model 给整条回复打分
    4. 用 critic 估计 value 并计算优势（GAE）
    5. 用 PPO 更新 actor 和 critic
    """
    for step, batch in enumerate(self.dataloader):
        # 生成样本
        samples = self.get_sample(batch)
        # 生成经验，包括 rollout 和 evaluation 阶段的 adv rewards 等等
        experience = self.get_experience(samples)
        torch.cuda.empty_cache()
        # 每个 prompt 都要进行多轮训练
        for epoch in range(self.config.ppo_epochs):
            self.train_step(experience)

定义: Actor rollout → ref/reward inference → critic inference → critic training → actor training
执行: verl 自动调度，分配 GPU，切换并行模式

# stage1 
if __name__ == '__main__':
    main()

# stage2
@hydra.main(config_path='config', config_name='ppo_trainer', version_base=None)
def main(config):
    run_ppo(config)

# stage3
def run_ppo(config, compute_score=None):
	# ...
    ray.get(main_task.remote(config, compute_score))

# stage4
@ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
def main_task(config, compute_score=None):
	global_pool_id = "global_pool"
	resource_pool_spec = {
	  global_pool_id: ([config.trainer.n_gpus_per_node] * config.trainer.nnodes),
	}
	mapping = {
	  Role.ActorRollout: global_pool_id, Role.Critic: global_pool_id,
	  Role.RefPolicy: global_pool_id, Role.RewardModel: global_pool_id,
	}
	resource_pool_manager = ResourcePoolManager(
	  resource_pool_spec=resource_pool_spec, mapping=mapping)
	# ...
	trainer = RayPPOTrainer(
		config=config, 
	    resource_pool_manager=resource_pool_manager, # ...
	)
	trainer.init_workers()
	trainer.fit()

def init_workers(self):
	all_wg = {}
	self.wg_dicts = []
	for resource_pool, class_dict in self.resource_pool_to_cls.items():
	    worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
	    wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)
	    spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
	    all_wg.update(spawn_wg)
	    # keep the referece of WorkerDict to support ray >= 2.31. Ref: https://github.com/ray-project/ray/pull/45699
	    self.wg_dicts.append(wg_dict)

ray_trainer.py / fit()          ← 第4层：single-controller，只管调度逻辑
    ↓ 调用
RayWorkerGroup                  ← 第3层：multi-controller，类似 torchrun，把方法分布式广播给组内所有 Worker
    ↓ 调用
fsdp_workers.py                 ← 第2层：ActorRolloutRefWorker，调度层
  内部持有 ↓                          把 rollout、actor training、ref 三个角色合并管理
dp_actor.py                     ← 第1层：DataParallelPPOActor，执行层
                                      真正跑前向/反向传播的地方

fit()
├── 初始化（加载checkpoint、验证）
└── 训练主循环（epoch → batch）
    ├── 生成阶段（Rollout）
    ├── 奖励计算阶段
    ├── 对数概率计算阶段
    ├── 优势估计阶段（Advantage）
    ├── Critic/Actor更新阶段
    └── 日志/验证/保存

# 1. 准备数据，只取 prompt 部分去掉 response
batch = DataProto.from_single_dict(batch_dict)
gen_batch = self._get_gen_batch(batch)

# 2. rollout 生成回答
#    底层通过 Ray 进程间通信把 DataProto 传给 rollout worker，计算完再传回
gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch)

# 3. 把生成结果合并回 batch
batch = batch.union(gen_batch_output)

# 4. 奖励 → KL 惩罚 → 优势估计
reward_tensor = extract_reward(batch)
batch = apply_kl_penalty(batch)
batch = compute_advantage(batch)    # GAE / GRPO 等

# 5. 更新网络
critic_output = self._update_critic(batch)
actor_output  = self._update_actor(batch)
self.checkpoint_manager.update_weights(...)  # 同步最新权重到 rollout worker

ActorRolloutRefWorker
├── Actor    → 用 PPO 更新模型参数（训练）
├── Rollout  → 用当前模型生成回答（采样）
└── Ref      → 用原始模型算 KL 散度用的参考概率（不训练）

async def rollout_mode(self):
    # 把 FSDP 分片的参数重新聚合成完整参数，推送给推理引擎
    params = self.actor_module_fsdp.state_dict()
    await self.rollout.update_weights(params)
    await self.rollout.resume(tags=["kv_cache"])   # 恢复 KV cache，准备生成

def generate_sequences(self, prompts):
    loop.run_until_complete(self.rollout_mode())           # 切换到推理模式
    output = self.rollout.generate_sequences(prompts)      # vLLM/SGLang 生成
    loop.run_until_complete(self.trainer_mode())           # 切回训练模式
    return output

DataProto (全量 batch)
    ↓  split
mini_batch × ppo_epochs        ← PPO 多轮更新
    ↓  split  
micro_batch × gradient_accum   ← 梯度累积（省显存）
    ↓
_forward_micro_batch()          ← 实际 forward
    ↓
loss.backward()
    ↓
_optimizer_step()               ← clip grad + step

for epoch in ppo_epochs:
    for mini_batch in mini_batches:
        optimizer.zero_grad()
        for micro_batch in micro_batches:           # 梯度累积
            log_prob = forward(micro_batch)
            pg_loss = policy_loss_fn(old_log_prob, log_prob, advantages)
            loss = pg_loss - entropy * coeff + kl_loss * coeff
            (loss * scale).backward()
        clip_grad_norm + optimizer.step()

def _init_with_resource_pool(self, resource_pool, ray_cls_with_init):
  # ...
  rank = -1
  for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)): # Node
    for local_rank in range(local_world_size): # GPU
      rank += 1
      env_vars = {
        'WORLD_SIZE': str(world_size), 'RANK': str(rank), # More env vars ...
      }
      ray_cls_with_init.update_options(
        {'runtime_env': {'env_vars': env_vars}})
      # ...
      worker = ray_cls_with_init(placement_group=pg,
                                 placement_group_bundle_idx=local_rank)
      self._workers.append(worker)
  # ...

def register(
	dispatch_mode=Dispatch.ALL_TO_ALL,
    execute_mode=Execute.ALL,
    blocking=True,
    materialize_futures=True
):

    def decorator(func):
        @wraps(func)
        def inner(*args, **kwargs):
            if materialize_futures:
                args, kwargs = _materialize_futures(*args, **kwargs)
            return func(*args, **kwargs)

        attrs = {
            'dispatch_mode':  dispatch_mode,
            'execute_mode':   execute_mode,
            'blocking':       blocking,
        }
        setattr(inner, MAGIC_ATTR, attrs)
        return inner

    return decorator

predefined_dispatch_mode_fn = {
    Dispatch.ONE_TO_ALL: {
        'dispatch_fn': dispatch_one_to_all,
        'collect_fn': collect_all_to_all,
    },
    Dispatch.ALL_TO_ALL: {
        'dispatch_fn': dispatch_all_to_all,
        'collect_fn': collect_all_to_all,
    },	
    Dispatch.DP_COMPUTE_PROTO: {
        'dispatch_fn': dispatch_dp_compute_data_proto,
        'collect_fn': collect_dp_compute_data_proto
    },
    # ...
}

# dispatch：把 DataProto 按 world_size 均匀切分，每个 Worker 拿到一片
def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(
        worker_group.world_size, *args, **kwargs
    )
    return splitted_args, splitted_kwargs

# collect：把所有 Worker 返回的 DataProto 片段拼回完整数据
def collect_dp_compute_data_proto(worker_group, output):
    return _concat_data_proto_or_future(output)

def _concat_data_proto_or_future(output):
    o = output[0]
    if isinstance(o, DataProto):
        return DataProto.concat(output)
    elif isinstance(o, ray.ObjectRef):
        return DataProtoFuture.concat(output)
    else:
        raise NotImplementedError

import ray
import math
from verl.single_controller.base import Worker
from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
from verl.single_controller.base.decorator import register, Dispatch, MAGIC_ATTR

# ─── Step 1: 定义自定义 dispatch/collect 函数 ────────────────────────────────

def dispatch_list_split(worker_group, *args, **kwargs):
    """把 prompts list 均匀切分成 world_size 份"""
    world_size = worker_group.world_size
    prompts = kwargs['prompts']
    
    chunk_size = math.ceil(len(prompts) / world_size)
    # 切成 world_size 个子列表，不足的 worker 分到空列表
    chunks = [
        prompts[i * chunk_size : (i + 1) * chunk_size]
        for i in range(world_size)
    ]
    # 返回的 kwargs 里 prompts 是个长度==world_size 的 list
    # execute_all_async 会把 chunks[i] 发给 worker[i]
    return args, {**kwargs, 'prompts': chunks}


def collect_list_concat(worker_group, output):
    """把各 Worker 返回的子列表拼回完整列表"""
    # output 是 [result_from_worker0, result_from_worker1, ...]
    result = []
    for sublist in output:
        result.extend(sublist)
    return result


# ─── Step 2: 把自定义逻辑注册到 Dispatch 映射表 ────────────────────────────────

# 扩展 predefined_dispatch_mode_fn（verl 内部的 dict）
from verl.single_controller.base.decorator import predefined_dispatch_mode_fn

# 用一个字符串 key 或自定义 Enum 值都可以，这里用字符串简单演示
DISPATCH_LIST_SPLIT = "LIST_SPLIT"

predefined_dispatch_mode_fn[DISPATCH_LIST_SPLIT] = {
    'dispatch_fn': dispatch_list_split,
    'collect_fn':  collect_list_concat,
}


# ─── Step 3: 在 Worker 里使用 ────────────────────────────────────────────────

@ray.remote
class InferWorker(Worker):
    def __init__(self):
        super().__init__()

	@register(dispatch_mode=Dispatch.ONE_TO_ALL):
	def load_model(self, model_name: str):
		self.model = AutoModelForCausalLM.from_pretrained(
			model_name,
			dtype="auto",
			device=f"cuda-{self.rank}"
		)
		self.tokenizer = AutoTokenizer.from_pretrained(model_name)
		
		return self.model.device, id(self.model)

    @register(dispatch_mode=DISPATCH_LIST_SPLIT)
    def infer(self, prompts: list[str]) -> list[str]:
        # 这里 prompts 已经是切分后的子列表，单卡逻辑
        print(f"[rank {self.rank}] received {len(prompts)} prompts: {prompts}")
        results = self.model.generate(prompts)
        return results


# ─── Step 4: 启动并调用 ─────────────────────────────────────────────────────

ray.init()

resource_pool   = RayResourcePool([2], use_gpu=True)  # 2 个 Worker
class_with_args = RayClassWithInitArgs(cls=InferWorker)
worker_group    = RayWorkerGroup(resource_pool, class_with_args)

model_info = worker_group.load_model("qwen3.5-2b")
print(model_info)

prompts = ["问题A", "问题B", "问题C", "问题D", "问题E"]
results = worker_group.infer(prompts=prompts)

print(results)
# ['[rank0] response to: 问题A', '[rank0] response to: 问题B', '[rank0] response to: 问题C',
#  '[rank1] response to: 问题D', '[rank1] response to: 问题E']

def execute_all_async(self, method_name: str, *args, **kwargs):
    # 这里我们假设，如果 args 和 kwargs 里面所有的参数都是 list，且所有的 list 长度都与 len(self._workers) 一致的话，我们会把
    # list 中的每一个分别发到对应的 worker 上去
    # print(f"execute_all_async: method {method_name}({args}, {kwargs})")
    length = len(self._workers)
    if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):
        if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):
            # print(f"splitting args and kwargs into {length} shards")
            result = []
            for i in range(length):
                sliced_args = tuple(arg[i] for arg in args)
                sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
                remote_call = getattr(self._workers[i], method_name)
                result.append(remote_call.remote(*sliced_args, **sliced_kwargs))
            return result
    
    return [getattr(worker, method_name).remote(*args, **kwargs) for worker in self._workers]

execute_all_async("generate_sequences")  # rollout 阶段
execute_all_async("compute_log_prob")    # 计算 logprob 阶段
execute_all_async("update_actor")        # 训练阶段

def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):

    def func(*args, **kwargs):
        # Step 1: 切分数据
        args, kwargs = dispatch_fn(self, *args, **kwargs)
        # Step 2: 分发给所有 Worker 异步执行
        output = execute_fn(method_name, *args, **kwargs)
        # Step 3: 阻塞等待（可选）
        if blocking:
            output = ray.get(output)
        # Step 4: 聚合结果
        output = collect_fn(self, output)
        return output

    return func

Driver（单进程）
  └── single-controller：顺序调用各 WorkerGroup 的方法
        └── WorkerGroup（multi-controller）：SPMD 调度
              └── dispatch → execute（多 Worker 并行）→ collect

actor_wg.update_actor(data)

{
	"data_source": used to chose reward function,
    "prompt": [{"role": ..., "content": ...}],
    "reward_model": {
        "style": "rule" or "reward",
        "ground_truth": ...
    },
    "extra_info": a dict containing extra information
}

def make_map_fn(split):

    def process_fn(example, idx):
        question_raw = example.pop('question')
        question = question_raw + ' ' + instruction_following
        answer_raw = example.pop('answer')
        solution = extract_solution(answer_raw)
        data = {
            "data_source": data_source,
            "prompt": [{
                "role": "user",
                "content": question,
            }],
            "ability": "math",
            "reward_model": {
                "style": "rule",
                "ground_truth": solution
            },
            "extra_info": {
                'split': split,
                'index': idx,
                'answer': answer_raw,
                "question": question_raw,
            }
        }
        return data

    return process_fn

train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)

train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet"))
test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet"))

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$HOME/data/geo3k/train.parquet \
    data.val_files=$HOME/data/geo3k/test.parquet \
    //...

custom_reward_function:
  path: /path/to/my_reward.py   # 你的文件路径
  name: my_reward_fn            # 函数名，如果叫 compute_score 可以不填
reward_model:
  reward_manager: naive

def compute_score(data_source, solution_str, ground_truth, extra_info=None):
    if data_source == "openai/gsm8k":
        return gsm8k_score(solution_str, ground_truth)
    elif data_source == "lighteval/MATH":
        return math_score(solution_str, ground_truth)
    else:
        raise NotImplementedError(f"Unknown data_source: {data_source}")

class NaiveRewardManager:
    def __call__(self, data: DataProto):
        for i in range(len(data)):
            data_item = data[i]
            sequences_str = self.tokenizer.decode(
                torch.cat((valid_prompt_ids, valid_response_ids))
            )
            score = self.compute_score(
                data_source=data_item.non_tensor_batch['data_source'],
                solution_str=sequences_str,
                ground_truth=data_item.non_tensor_batch['reward_model']['ground_truth'],
            )
            reward_tensor[i, valid_response_length - 1] = score
        return reward_tensor

class PrimeRewardManager:
    def __call__(self, data: DataProto):
        try:
            scores = asyncio.run(
                parallel_compute_score_async(
                    self.compute_score,
                    sequences_str, ground_truth, data_sources,
                    num_processes=64,
                )
            )
        except asyncio.TimeoutError:
            scores = [0.0] * len(sequences_str)
        # 写回 reward tensor ...
        return reward_tensor

def compute_score(data_sources, solution_strs, ground_truths, extra_infos=None):
	return [0.0] * len(solution_strs)

def compute_score(data_source, solution_str, ground_truth, extra_info=None):
    format_score      = check_format(solution_str)
    correctness_score = check_answer(solution_str, ground_truth)
    length_penalty    = -0.01 * max(0, len(solution_str) - 1000)
    return 0.2 * format_score + 0.7 * correctness_score + 0.1 * length_penalty

class MultiRewardManager:
    def __init__(self, tokenizer, num_examine, compute_score=None):
        self.tokenizer = tokenizer
        self.reward_fns = [
            (compute_format_score,      0.2),
            (compute_correctness_score, 0.7),
            (compute_length_score,      0.1),
        ]

    def __call__(self, data: DataProto):
        # decode、提取 ground_truth ...
        total = sum(
            w * fn(data_source, solution_str, ground_truth, extra_info)
            for fn, w in self.reward_fns
        )
        # 写回 reward tensor ...
        return reward_tensor

def compute_policy_loss_xxx(
    old_log_prob,      # (bsz, response_len) rollout 时采样的 log prob
    log_prob,          # (bsz, response_len) 当前 policy 的 log prob  
    advantages,        # (bsz, response_len) GAE 优势估计
    response_mask,     # (bsz, response_len) 哪些位置是有效 response
    loss_agg_mode,     # "token_mean" / "seq_mean" / ...
    config,
    **kwargs
) -> (loss_scalar, metrics_dict)

actor_rollout_ref:
  actor:
    policy_loss:
      loss_mode: compute_policy_loss_xxx
    entropy_coeff: 0.001
    use_kl_loss: false

def compute_custom_advantage(  
    token_level_rewards: torch.Tensor,  
    response_mask: torch.Tensor,  
    index: np.ndarray = None,  
    config: Optional[AlgoConfig] = None,  
    **kwargs  
) -> tuple[torch.Tensor, torch.Tensor]:  
    """  
    计算自定义优势  
      
    Returns:  
        advantages: 优势值 (bs, response_length)  
        returns: 回报值 (bs, response_length)  
    """  
    # 实现你的优势计算逻辑  
    return advantages, returns

from verl.trainer.ppo.core_algos import register_adv_est, AdvantageEstimator  
  
@register_adv_est("custom_advantage")
def compute_custom_advantage(...):
    pass

algorithm:  
  adv_estimator: custom_advantage

select_keys = [
    "responses", "response_mask", "input_ids",
    "attention_mask", "position_ids",
    "old_log_probs", "advantages",
    "token_level_rewards",   # ← 新增字段
]

# 然后在 forward 循环里取出来传给 loss fn
token_rewards = model_inputs.get("token_level_rewards", None)

pg_loss, pg_metrics = policy_loss_fn(
    old_log_prob=old_log_prob,
    log_prob=log_prob,
    advantages=advantages,
    response_mask=response_mask,
    loss_agg_mode=loss_agg_mode,
    config=self.config,
    token_level_rewards=token_rewards,  # ← 透传
)

def train(self):
    for epochs in range(self.config.epochs):
        for batch in self.dataloader: 
            exp = self.get_experience(batch["prompts"])
            
            for _ in range(self.config.ppo_epochs):
                for mini_batch in split(exp, self.config.ppo_mini_batch_size):
	                for micro_batch in split(mini_batch, micro_batch_size_per_gpu * dp_size):
	                    self.optimizer.zero_grad()
	                    loss = self.compute_loss(micro_batch)
	                    loss.backward()
	                    self.optimizer.step()

                    ┌─────────────────────────────────────────────┐
                    │              ToolAgentLoop.run()            │
                    └─────────────────────────────────────────────┘
                                          │
                              ┌───────────▼───────────┐
                              │    AgentState.PENDING │  ← 准备 prompt
                              └───────────┬───────────┘
                                          │
                           ┌──────────────▼──────────────┐
                           │   AgentState.GENERATING     │  ← LLM 推理
                           └──────────────┬──────────────┘
                                          │
                        ┌─────────────────┼───────────────────┐
                        │                 │                   │
               有 tool_calls        有 interaction        无任何后续
                        │                 │                   │
         ┌──────────────▼──────┐  ┌───────▼────────┐  ┌──────▼──────────┐
         │ PROCESSING_TOOLS    │  │  INTERACTING   │  │   TERMINATED    │
         │ (执行工具, 拼 token) │  │ (获取用户输入)  │  │   (输出结果)     │
         └──────────────┬──────┘  └───────┬────────┘  └─────────────────┘
                        │                 │
                        └────────┬────────┘
                                 │ 回到 GENERATING
                                 ▼
                      (直到 TERMINATED 退出循环)

class AgentData:
    # ── 输入数据 ──────────────────────────────
    messages: list[dict]        # 对话历史（role/content 格式）
    image_data: list[Image]     # 多模态图像
    video_data: list[tuple]     # 多模态视频
    tools_kwargs: dict          # 工具初始化参数
 
    # ── 训练关键字段 ───────────────────────────
    prompt_ids: list[int]       # 完整的 token id 序列（prompt + 所有 response）
    response_ids: list[int]     # 当前轮次的 response token ids
    response_mask: list[int]    # 1=模型生成的 token, 0=工具/用户输入的 token
    response_logprobs: list[float]  # 对应 logprob，用于 PPO 等算法
 
    # ── 统计与奖励 ─────────────────────────────
    turn_scores: list[float]    # 每轮 interaction 产生的奖励
    tool_rewards: list[float]   # 每次工具调用产生的奖励
    user_turns: int
    assistant_turns: int
 
    # ── 临时状态 ───────────────────────────────
    tool_calls: list[FunctionCall]  # 当前轮次解析出的工具调用
    extra_fields: dict              # 自定义扩展字段

while state != AgentState.TERMINATED:
    if state == AgentState.PENDING:
        state = await self._handle_pending_state(agent_data, sampling_params)
    elif state == AgentState.GENERATING:
        state = await self._handle_generating_state(agent_data, sampling_params)
    elif state == AgentState.PROCESSING_TOOLS:
        state = await self._handle_processing_tools_state(agent_data)
    elif state == AgentState.INTERACTING:
        state = await self._handle_interacting_state(agent_data)
    else:
        logger.error(f"Invalid state: {state}")
        state = AgentState.TERMINATED

async def _handle_pending_state(self, agent_data, sampling_params):
    prompt_ids = await self.apply_chat_template(
        agent_data.messages,
        tools=schemas,      # 把工具 schema 嵌入 prompt
        images=...,
        videos=...,
    )
    agent_data.prompt_ids = prompt_ids
    return AgentState.GENERATING

async def _handle_generating_state(self, agent_data, sampling_params):
    # 1. 调用 LLM 生成
    output = await self.server_manager.generate(
        prompt_ids=agent_data.prompt_ids, ...
    )
 
    # 2. 累积 token 序列
    agent_data.response_ids = output.token_ids
    agent_data.prompt_ids += agent_data.response_ids    # prompt 不断增长
    agent_data.response_mask += [1] * len(agent_data.response_ids)  # 模型输出 mask=1
 
    # 3. 检查终止条件
    if len(agent_data.response_mask) >= self.response_length: return TERMINATED
    if agent_data.assistant_turns >= self.max_assistant_turns: return TERMINATED
 
    # 4. 解析 tool_calls
    _, agent_data.tool_calls = await self.tool_parser.extract_tool_calls(...)
 
    # 5. 决定下一状态
    if agent_data.tool_calls:      return PROCESSING_TOOLS
    elif self.interaction_config:  return INTERACTING
    else:                          return TERMINATED

tool_calls (并行执行，上限 max_parallel_calls)
    │
    ▼
asyncio.gather(*tasks)    ← 并行调用所有工具
    │
    ▼
构造 tool 消息（text / 多模态）
    │
    ▼
apply_chat_template(tool_messages)  ← 转为 token ids
    │
    ▼
prompt_ids += tool_response_ids
response_mask += [0] * len(tool_response_ids)   ← 工具输出 mask=0
    │
    ▼
return GENERATING   ← 继续让模型生成

async def _call_tool(self, tool_call, tools_kwargs, agent_data):
    try:
        instance_id, _ = await tool.create(...)
        response, reward, res = await tool.execute(instance_id, tool_args, agent_data=agent_data)
    except Exception as e:
        return ToolResponse(text=f"Error when executing tool: {e}"), 0.0, {}
    finally:
        await tool.release(instance_id)   # 确保资源释放

if len(tool_response_text) > self.max_tool_response_length:
    # 三种截断模式：left / right / middle

async def _handle_interacting_state(self, agent_data):
    should_terminate, response, reward, metrics = \
        await agent_data.interaction.generate_response(...)
 
    # 将用户回复追加进去，mask=0（非模型生成）
    response_ids = await self.apply_chat_template([{"role": "user", "content": response}])
    agent_data.prompt_ids += response_ids
    agent_data.response_mask += [0] * len(response_ids)
 
    if should_terminate: return TERMINATED
    else: return GENERATING

tool_list = initialize_tools_from_config(tool_config_path)
self.tools = {tool.name: tool for tool in tool_list}
self.tool_schemas = [tool.tool_schema.model_dump(...) for tool in tool_list]

tool_selection = extra_info.get("tool_selection")
if tool_selection:
    agent_data._active_tools = {name: self.tools[name] for name in tool_selection}
    agent_data._active_tool_schemas = [...]
else:
    agent_data._active_tools = self.tools   # 使用全部工具

actor_rollout_ref:
  rollout:
    multi_turn:
      format: hermes

# verl/experimental/agent_loop/tool_parser.py
import regex, json
from verl.experimental.agent_loop.tool_parser import ToolParser, FunctionCall

@ToolParser.register("my_format")
class MyModelToolParser(ToolParser):
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        # 根据你模型实际输出的格式定义正则
        self.pattern = regex.compile(
            r"<tool_call>(.*?)</tool_call>", regex.DOTALL
        )

    async def extract_tool_calls(
        self, response_ids: list[int]
    ) -> list[FunctionCall]:
        text = self.tokenizer.decode(response_ids, skip_special_tokens=False)
        matches = self.pattern.findall(text)
        calls = []
        for m in matches:
            try:
                obj = json.loads(m)
                calls.append(FunctionCall(
                    name=obj["name"],
                    arguments=json.dumps(obj["arguments"], ensure_ascii=False)
                ))
            except Exception:
                pass   # 格式解析失败，忽略该次调用
        return calls

format: my_format   # 对应 @ToolParser.register 的 key

AgentLoopOutput(
    prompt_ids   = prompt_ids,                         # 纯 prompt 部分
    response_ids = response_ids[:self.response_length], # 截断到最大长度
    response_mask= response_mask[:self.response_length],
    response_logprobs = ...,
    num_turns    = user_turns + assistant_turns + 1,
    metrics      = {...},
    extra_fields = {
        "turn_scores": [...],   # 每轮 interaction 奖励
        "tool_rewards": [...]   # 每次工具调用奖励
    }
)

from verl.tools.base_tool import BaseTool, ToolResponse
from typing import Tuple, Dict, Any

class MySearchTool(BaseTool):

    def get_openai_tool_schema(self) -> dict:
        """
        返回 OpenAI Function Call 格式的工具描述，
        VERL 用它拼接到推理请求的 tools 字段里，告诉模型有哪些工具可用。
        """
        return {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "搜索互联网，返回最多10条相关文档，每条包含 title 和 content",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "搜索关键词"
                        }
                    },
                    "required": ["query"]
                }
            }
        }

    async def execute(self, name: str, arguments: dict[str, Any]) -> ToolResult:
        """
        工具的实际执行逻辑。
        AgentLoop 解析出 Function Call 后，调用此方法拿到工具返回值，
        再把结果拼回对话的 tool_response 里继续推理。
        
        name:      模型调用的函数名，比如 "web_search"
        arguments: 模型传入的参数，比如 {"query": "北京天气"}
        返回:      ToolResult，包含返回给模型的文本内容
        """
        query = arguments.get("query", "")
        # 替换成你真实的检索服务
        raw = await self._call_retrieval_service(query)
        return ToolResult(content=raw, is_error=False)

    def calc_reward(self) -> float:
        """
        （可选）基于工具调用状态计算过程奖励，例如：
        - 是否成功调用了工具（格式是否正确）
        - 工具调用次数是否合理（防止滥用）
        - 工具返回结果是否有效
        若不需要过程奖励，直接 return 0.0 即可。
        """
        return 0.0

# search_tool_config.yaml
tools:
  - class_name: verl.tools.search_tool.SearchTool  # 或你的自定义类
    config:
      retrieval_service_url: http://127.0.0.1:8000/retrieve
      num_workers: 120
      rate_limit: 120
      timeout: 30

actor_rollout_ref:  
  rollout:
	name: vllm
    mode: async  # 异步模式，避免GPU空闲  
    multi_turn:  
      enable: True  
      max_user_turns: 3  # 最大用户轮次  
      max_assistant_turns: 3  # 最大助手轮次 
    tool_kwargs:
      tools_config_file: ./config/tool_config/search_tool_config.yaml

{
	"data_source": used to chose reward function,
    "prompt": [{"role": ..., "content": ...}],
    "reward_model": {
        "style": "rule" or "reward",
        "ground_truth": ...
    },
    "extra_info": a dict containing extra information,
    "agent_name": "tool_agent"
}

from verl.experimental.agent_loop.base import AgentLoopBase, AgentLoopOutput, register
from verl.experimental.agent_loop.tool_parser import ToolParser
import aiohttp, json

@register("my_search_agent")
class MySearchAgentLoop(AgentLoopBase):

    async def run(self, sampling_params: dict, **kwargs) -> AgentLoopOutput:
        """
        必须实现的唯一接口。
        负责完整执行一条样本的 Multi-Turn 推理+工具调用，
        返回 AgentLoopOutput，其中 response_mask 标记哪些 token 参与 loss 计算。
        
        VERL 框架只关心你最终返回什么 token，
        中间怎么调工具、循环几轮，完全由你控制。
        """
        messages = list(kwargs["raw_prompt"])   # 原始对话历史
        max_turns = self.config.get("max_turns", 5)

        all_response_ids = []
        all_response_mask = []   # 1=参与loss，0=不参与（工具返回内容不算模型生成）

        for turn in range(max_turns):
            # 1. 调推理引擎，生成一段回复
            output = await self.llm_engine.chat(
                messages=messages,
                sampling_params=sampling_params,
                tools=[self.tool_schema]
            )
            response_text = output.text
            response_ids  = output.token_ids

            # 模型自己生成的 token，全部参与 loss
            all_response_ids  += response_ids
            all_response_mask += [1] * len(response_ids)

            # 2. 检测是否有 Function Call
            tool_calls = await self.tool_parser.extract_tool_calls(response_ids)

            if not tool_calls:
                # 没有工具调用，模型直接给出答案，结束循环
                messages.append({"role": "assistant", "content": response_text})
                break

            # 3. 执行工具调用
            messages.append({"role": "assistant", "content": response_text})
            for fc in tool_calls:
                tool_result = await self._call_search(fc.arguments["query"])
                tool_response_ids = self.tokenizer.encode(tool_result)

                # 工具返回内容不是模型生成的，mask=0，不参与 loss
                all_response_ids  += tool_response_ids
                all_response_mask += [0] * len(tool_response_ids)

                messages.append({
                    "role": "tool",
                    "name": fc.name,
                    "content": tool_result
                })

        return AgentLoopOutput(
            response_ids=all_response_ids[: self.response_length],
            response_mask=all_response_mask[: self.response_length],
        )

    async def _call_search(self, query: str) -> str:
        url = self.config["retrieval_url"]
        async with aiohttp.ClientSession() as s:
            async with s.post(url, json={"query": query}) as r:
                data = await r.json()
        return json.dumps(data, ensure_ascii=False)

data:
  return_raw_chat: True          # 必须，把原始 chat messages 传给 agent loop

actor_rollout_ref:
  rollout:
    mode: async                  # 必须，启用 server-based 异步 rollout
    name: sglang                 # 或 vllm
    
    agent_loop:
      # 告诉 verl 用哪个类处理 agent_name="my_search_agent" 的样本
      my_search_agent:
        class_name: my_search_agent_loop.MySearchAgentLoop
        config:
          retrieval_url: http://127.0.0.1:8000/retrieve
          max_turns: 5
      
      default:
        class_name: verl.agent_loop.SingleTurnAgentLoop
    
    # 并发控制
    agent_loop_kwargs:
      max_concurrent: 128

dataset = dataset.map(lambda x: {
    **x,
    "agent_name": "my_search_agent"  # 对应你注册的名字
})

class AgentLoopWorker:
	async def generate_sequence():
		# ...
		tasks = []
        for i in range(len(batch)):
            trace_this_sample = i in traced_indices
            kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items()}
            tasks.append(
                asyncio.create_task(
                    self._run_agent_loop(sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs)
                )
            )
        outputs = await asyncio.gather(*tasks)
        output = self._postprocess(
            outputs, input_non_tensor_batch=batch.non_tensor_batch, validate=batch.meta_info.get("validate", False)
        )
        return output

Rollout 群集: |--- 生成 Batch 1 --|--- 生成 Batch 2 --|--- 生成 Batch 3 --|
              (使用权重: θ0)       (使用权重: θ0)       (使用权重: θ1)
                    |                    |                    |
             [塞入 Queue 1]       [塞入 Queue 2]        [塞入 Queue 3]
                    |                    |                    |
Trainer 群集:       |-- 训练 Batch 1 --|-- 训练 Batch 2 --|-- 训练 Batch 3 --|
                       (θ0 -> θ1)           (θ1 -> θ2)           (θ2 -> θ3)
                            |                    |
                            v                    v
                      [异步广播 θ1]         [异步广播 θ2]