训练初始化

训练从任务集开始，任务集是一类共享输入格式和奖励的任务。格式器将每个任务转换为智能体提示，奖励函数对智能体的联合响应进行评分。

任务集#

任务集是 Hugging Face、Dataset 或 IterableDataset 的行，每行是一个字典，代表任务项。默认格式器读取 prompt 字段；其他字段可以包含目标、标签、测试用例或其他奖励函数所需的其他信息。Trainers 将原始任务项通过 batch_items 传递给奖励函数。

from datasets import Dataset

tasks = Dataset.from_dict({
    "prompt": [
        "Describe a city in the clouds.",
        "Invent a new holiday.",
        "Write a bedtime story for a dragon.",
        "Explain how teleportation might work.",
        "Tell a story about dinosaurs.",
        "Describe a world without electricity.",
        "Create a superhero with a unique power.",
        "Write a scene where the moon talks.",
        "Invent a new type of fruit.",
        "Design a playground on Mars.",
    ],
    "target_min": [2.0] * 10,
    "target_max": [3.0] * 10,
})

split = tasks.train_test_split(test_size=0.2, seed=42)
train_dataset = split["train"]
eval_dataset = split["test"]

这里，八个任务项用于训练，两个用于评估。seed 确保分割的可重复性；它不为模型生成设置随机种子。现有的训练和评估分割可以直接使用。

当前的训练数据加载器保留任务顺序。用户可以使用 train_dataset.shuffle(seed=42) 在创建 trainer 之前对任务项进行洗牌。这会产生固定的顺序，而无需重新洗牌每个 epoch。

任务提示#

一个格式器可以共享提示格式给多个智能体，或者 actor 中格式器的顺序可以分配不同的角色。仅用于格式器的字段进入模型提示；奖励函数仍然可以读取完整的任务项。

def brief_formatter(item):
    return f"Write a short response to this task:\n{item['prompt']}"

def detailed_formatter(item):
    return f"Write a detailed response to this task:\n{item['prompt']}"

formatters = [brief_formatter, detailed_formatter]

对于多轮任务，格式器还可以接受 external_prompts。过渡接口在多轮交互中描述。

奖励#

基于规则的奖励和 LLM 奖励模型都通过 reward_func 传递给 trainers。奖励函数接收来自每个智能体的完成列表，顺序为 actor。如果它接受 batch_items 关键词，trainers 也会传递任务项。它返回一个联合响应的奖励。 trainers 用于此处的函数，一次处理一个任务项，该项可用为 batch_items[0]。

reward_processor 在训练前会转换每个返回的奖励；例如，RewardProcessors.scale(factor=10.0) 从 comlrl.utils.reward_processor 将奖励乘以 10。默认情况下，奖励保持不变。

基于规则的奖励#

基于规则的奖励直接从响应和任务项中计算分数，例如通过任务检查、验证器或测试框架。以下函数对长度比任务进行评分：

import math

def reward_func(brief_completions, detailed_completions, *, batch_items):
    item = batch_items[0]
    rewards = []
    for brief, detailed in zip(brief_completions, detailed_completions):
        if not brief:
            rewards.append(0.0)
            continue
        ratio = len(detailed) / len(brief)
        distance = max(item["target_min"] - ratio, ratio - item["target_max"], 0.0)
        rewards.append(math.exp(-distance))
    return rewards

此函数返回一个与对齐响应的共享奖励。任务对字符数量比进行评分；它不衡量写作质量，并且智能体独立生成。

用户可以在不加载模型的情况下检查任务集、格式器和奖励函数：

item = train_dataset[0]
assert item["prompt"] in brief_formatter(item)
assert reward_func(["Hi"], ["Hello"], batch_items=[item]) == [1.0]
assert reward_func([""], ["Hello"], batch_items=[item]) == [0.0]

基于 LLM 的奖励模型#

奖励函数也可以使用现有的 LLM 奖励模型对响应进行评分； trainers 只使用它返回的值。以下示例将奖励模型与单个评分分类头包装； reward_model_id 是占位符：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

reward_model_id = "your-org/your-reward-model"
reward_tokenizer = AutoTokenizer.from_pretrained(reward_model_id)
reward_model = AutoModelForSequenceClassification.from_pretrained(reward_model_id).eval()

@torch.no_grad()
def model_reward(brief_completions, detailed_completions, *, batch_items):
    task = batch_items[0]["prompt"]
    texts = [
        f"{task}\n\nBrief: {brief}\n\nDetailed: {detailed}"
        for brief, detailed in zip(brief_completions, detailed_completions)
    ]
    inputs = reward_tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
    return reward_model(**inputs).logits[:, 0].tolist()

# Pass it when constructing a trainer:
# reward_func=model_reward,

奖励模型可以在函数外部加载一次，并将其放置在具有可用内存的设备上； trainers 不会移动或更新它。要从偏好中学习奖励模型，请参阅 LLM 奖励模型。

训练和评估#

用户可以通过 train_dataset 和 eval_dataset 分别提供训练和评估任务项。评估使用训练格式器和奖励函数，但不更新模型权重。输出采用配置的生成设置，因此重复评估可能产生不同的分数。

from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer

config = MAGRPOConfig(
    num_agents=2,
    num_turns=1,
    num_train_epochs=1,
    num_generations=4,
    max_new_tokens=128,
    eval_interval=4,
    eval_num_samples=len(eval_dataset),
    eval_batch_size=1,
    logging_steps=4,
)

trainer = MAGRPOTrainer(
    agent_model="Qwen/Qwen2.5-0.5B",
    num_agents=2,
    formatters=formatters,
    reward_func=reward_func,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=config,
)

trainer.train()
metrics = trainer.evaluate(num_eval_samples=len(eval_dataset))
print(metrics)

对于 MAREINFORCE、IAC 和 MAAC 系列，eval_interval 计入训练任务项。评估在每个 epoch 的第一个任务项之前运行，然后按配置的间隔运行；0 仅启用手动评估。eval_num_samples 指定评估任务项的上限，从评估分割的开始处开始。将其设置为分割长度将评估所有任务项。eval_batch_size 控制评估数据加载器的批次大小，而不是每个智能体生成的响应数量。

MAGRPO、IAC 和 MAAC 在评估期间，每个智能体每轮生成一个响应。IAC 和 MAAC 使用其 rollout 路径，以及配置的轮次结构，但其 evaluate() 方法不接受样本计数参数：

# For an existing IACTrainer or MAACTrainer:
ac_trainer.args.eval_num_samples = len(eval_dataset)
metrics = ac_trainer.evaluate()

如果没有 eval_dataset，evaluate() 返回一个空字典。对于 PbRL trainers，手动 evaluate(num_eval_samples=...) 也是可用的；自动评估时间表遵循偏好更新或 actor-训练阶段，而不是上述任务项时间表。

日志#

CoMLRL 支持在 Weights & Biases 中记录训练曲线、任务特定指标和本地评估记录。示例使用 trainer 和上述评估分割。

Weights & Biases#

用户可以通过在创建 trainer 时提供 wandb_config 来在 Weights & Biases 中记录训练曲线：

wandb_config = {
    "project": "comlrl",
    "name": "length-ratio-magrpo",
    "tags": ["length-ratio", "magrpo"],
    "config_sections": {
        "dataset": {"name": "length-ratio", "split_seed": 42},
        "output": {"verbose": True},
    },
}

# Add this keyword when constructing MAGRPOTrainer:
# wandb_config=wandb_config,

W&B可以在用户环境中进行身份验证，在启用日志记录之前。忽略wandb_config或将其设置为None会禁用W&B。支持的运行设置包括project、entity、name（或run_name）、tags和dir。 config_sections字段将任务和实验设置添加到记录的配置中；config_sections.output.verbose控制trainer进度输出。

训练指标#

Trainers默认记录以下指标，每次轮次k（从1开始）。自定义指标可以将其与任务特定值一起添加。

指标	含义
`turn_k/reward_mean`	训练rollouts在轮次`k`的平均奖励
`turn_k/expected_return`	训练rollouts从轮次`k`的平均回报
`eval/turn_k/reward_mean`	评估rollouts在轮次`k`的平均奖励
`eval/turn_k/expected_return`	评估rollouts从轮次`k`的平均回报

启用参考KL时，reference_kl_mean和reference_kl_penalty_mean也会为每个轮次记录。训练指标记录在env_step上，即联合行动奖励评估的数量，该数量随着多个生成或轮次的增长而更快地增长。 logging_steps设置训练日志之间的最小env_step间隔；评估指标遵循eval_interval。

自定义指标#

用户可以通过提供 eval_logger 和 eval_aggregator 来为 MAGRPO 添加自定义评估指标。日志记录器接收到的完成结果包括 [agent][task item][turn] 顺序上的结果，以及原始 batch_items 和原始 prompts。聚合器返回的标量指标不包含 eval/ 前缀；trainer 会添加它。

以下示例测量了每个轮次中，在每个智能体中平均响应长度：

def eval_logger(*, agent_completions_turns, batch_items, prompts):
    return [
        {
            f"turn_{turn + 1}/response_chars": sum(
                len(agent[row][turn]) for agent in agent_completions_turns
            ) / len(agent_completions_turns)
            for turn in range(len(agent_completions_turns[0][row]))
        }
        for row in range(len(batch_items))
    ]

def eval_aggregator(rows, *, num_turns):
    return {
        key: sum(row[key] for row in rows) / len(rows)
        for turn in range(num_turns)
        for key in [f"turn_{turn + 1}/response_chars"]
    } if rows else {}

# Add both keywords when constructing MAGRPOTrainer:
# eval_logger=eval_logger,
# eval_aggregator=eval_aggregator,

生成的字典包括 eval/turn_1/response_chars。这些回调函数也接收任务-项目字段，允许在提示中计算任务特定的准确性或组指标，而无需参考答案。

IAC 和 MAAC 代替 metrics_callback(rollouts)，后者返回一个包含数字值的字典。每个 rollout 样本提供 agent_idx、prompt、completion 和 metadata 字段：

def metrics_callback(rollouts):
    if not rollouts:
        return {}
    return {
        "response_chars": sum(len(sample.completion) for sample in rollouts)
        / len(rollouts)
    }

# Add metrics_callback=metrics_callback to IACTrainer(...) or MAACTrainer(...).

这个回调函数在训练和评估期间运行。评估会将它的值平均到评估过的任务项目中，并添加 eval/；训练记录它的值在 rollout 收集时间，与标准 logging_steps 门分开。

本地指标#

即使 W&B 已禁用，评估也会返回指标字典。此字典可以本地保存以记录已完成的评估：

import json
from pathlib import Path

metrics = trainer.evaluate(num_eval_samples=len(eval_dataset))
Path("eval-metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")