训练实体 | CoMLRL

CoMLRL 支持训练 LLM 智能体，并使用基于 LLM 的 critics 或奖励模型进行训练。智能体和 critics 可以是来自 Hugging Face 或本地目录的同构或异构模型。用户可以通过 agent_model/critic_model，对于同构设置指定 Hugging Face 模型标识符，或者提供 agents/critics 列表，对于异构设置。训练好的实体可以保存并加载到新的 trainer 中，以便进行进一步训练。

LLM 智能体#

同构智能体#

同构智能体通过将 agent_model 设置为单个模型标识符并留空 agents=None 来初始化。对于具有单独的 num_agents 构造函数参数的 trainers，该参数可以与匹配的配置值一起设置；仅配置本身不会覆盖它。

以下示例加载三个 Qwen/Qwen2.5-1.5B 智能体：

trainer = MAGRPOTrainer(
    agent_model="Qwen/Qwen2.5-1.5B",
    agents=None,
    num_agents=3,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(num_agents=3, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

异构智能体#

虽然同构智能体可以被提示在不同的角色下并行处理任务的不同部分，但异构智能体可以贡献不同模型的专业能力。

用户可以通过 agents 供给模型标识符或已加载的模型对象，并使用 tokenizer 提供的兼容 tokenizers。使用 agents 提供的 agent_model 可以保持 None。如果同时提供，它们的模型名称必须匹配，并且它们的列表长度必须与智能体的数量匹配；否则，训练会引发错误。

以下示例加载 Qwen/Qwen2.5-Coder-3B 和 Qwen/Qwen2.5-Coder-7B 智能体:

trainer = MAGRPOTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(num_agents=2, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

LLM Critics#

诸如独立 Actor-Critic (IAC) 和多智能体 Actor-Critic (MAAC) 算法，使用 LLM critics。这些 critics 的加载方式取决于算法和 use_separate_critic 设置。 MARLHF trainers 与 rl_algorithm="maac" 或 "iac" 接受相同的 critic_model 和 critics 参数，或者在配置中使用 critic_model_name；否则，它们使用 actor 的模型标识符。

集中式 Critic#

可以通过 critic_model 或一个包含单个元素的 critics 列表加载一个单独的集中式 critic。它必须是一个受支持的因果语言模型：trainer 添加一个 value head 并使用第一个 actor 的 tokenizer 对联合输入进行编码，因此 tokenizer 必须与 critic 兼容。

以下示例将 Qwen/Qwen2.5-Coder-3B 和 Qwen/Qwen2.5-Coder-1.5B actors 与一个集中式 Qwen/Qwen2.5-Coder-7B critic 配对:

trainer = MAACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-1.5B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-1.5B"),
    ],
    critic_model="Qwen/Qwen2.5-Coder-7B",
    critics=None,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAACConfig(num_agents=2, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

独立 Critics#

使用独立 critics, 每个智能体可以拥有自己的 critic. 使用 use_separate_critic=True, 用户可以提供一个 critic_model 标识符来加载每个智能体单独的副本，或者一个 critics 长度为 num_agents. 每个单独的 critic 接收 token IDs 来自它的 actor’s tokenizer 并且必须与之兼容. 与 use_separate_critic=False 配合，每个 actor 和 critic 共享一个 LLM backbone，并一起加载。在这种情况下，critic_model 和 critics 保持未设置 (null 或 None)。与智能体类似，提供 critic_model 和 critics 需要匹配的模型名称和长度；否则，训练会引发错误。

以下示例使用基于相同模型的分开 critics 加载 Qwen/Qwen2.5-Coder-3B 和 Qwen/Qwen2.5-Coder-7B actors:

trainer = IACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B"),
    ],
    critic_model=None,
    critics=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=IACConfig(
        num_agents=2,
        num_turns=1,
        use_separate_critic=True,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)

或者，每个 actor 可以与 critic 共享相同的模型:

trainer = IACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B"),
    ],
    critic_model=None,
    critics=None,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=IACConfig(
        num_agents=2,
        num_turns=1,
        use_separate_critic=False,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)

Trainers 使用 agents/critics 内部列表。 agent_model 和 critic_model 参数是用于同构设置的快捷方式，并且必须与任何提供的相应列表保持一致。

MAGRPO 自动加载 tokenizers 作为模型标识符字符串。 MAAC/IAC 自动加载 tokenizer 作为 agent_model，但如果提供 agents，则需要明确指定 tokenizer，即使它包含字符串。加载的模型对象需要一个兼容的共享 tokenizer 或 tokenizer 列表，与 actors 匹配。异构 actors 在每个使用自己的兼容 tokenizer 时，可能具有不同的词汇。此外，每个 critics 必须接受 actor 的 token IDs (actor 的 IDs 用于 MAAC)。

LLM 奖励模型#

在 PbRL 中，MARLHF 从偏好对中学习联合奖励模型，并将其固定为政策学习的奖励。它的 backbone 通过 reward_model_name 在 MARLHFConfig 中加载，从 Hugging Face 模型标识符或因果语言模型的本地目录中加载。它默认为 actor 的模型标识符，并在 actors 是没有标识符的模型对象时需要。trainer 加载 backbone 与自己的 tokenizer，在最后一个非填充 token 上添加标量头，并将模型放置在 reward_model_device (默认 "cuda:0"; None 使用 trainer 设备)。reward_freeze_backbone=True 仅训练头。

以下示例加载一个 Qwen/Qwen2.5-Coder-7B 奖励模型，用于两个 Qwen/Qwen2.5-Coder-3B 智能体：

from comlrl.trainers.preference import MARLHFConfig, MARLHFTrainer

trainer = MARLHFTrainer(
    agent_model="Qwen/Qwen2.5-Coder-3B",
    num_agents=2,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MARLHFConfig(
        num_agents=2,
        num_turns=1,
        reward_model_name="Qwen/Qwen2.5-Coder-7B",
        reward_model_device="cuda:1",
    ),
)

MARLHFIterTrainer 在每次奖励学习迭代前，从该来源重新加载奖励模型。 reward_func 的任务奖励仍然标记偏好对；训练初始化描述了如何传递基于规则或 LLM 的任务奖励。

保存和恢复训练#

保存模型#

用户可以在训练后调用 save_model(output_dir)，以明确保存 actors 的 Hugging Face 模型配置、权重和 tokenizers 到本地目录:

trainer.train()
trainer.save_model("./saved-models/run-1")

从保存的模型中训练#

可以通过向原始角色顺序提供 actors 的目录，从保存的两个 actor MAGRPO 设置继续训练。父 run-1 目录包含 actors；每个 agent_i 目录都是一个可加载的模型。

from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer

continued_trainer = MAGRPOTrainer(
    agents=[
        "./saved-models/run-1/agent_0",
        "./saved-models/run-1/agent_1",
    ],
    num_agents=2,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(
        num_agents=2,
        num_turns=1,
        num_train_epochs=3,
        agent_learning_rate=5e-6,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)
continued_trainer.train()
continued_trainer.save_model("./saved-models/run-2")

MAGRPO 从 actor 目录中自动加载 tokenizers。在这里，num_train_epochs=3 表示新训练运行中的三个 epochs。数据集、奖励函数、训练配置以及任何自定义格式器、奖励处理器或外部过渡必须再次提供；save_model 不会存储它们。

其他 trainers 支持相同的本地目录方法，前提是满足 tokenizer 和 actor 的计数要求。偏好 trainers 开始偏好收集和根据新运行的配置进行奖励学习。

IAC 和 MAAC 将 value heads 分别保存为 value_head.pt。通过构造函数加载 critic 或共享 actor 目录会恢复语言模型 backbone 并创建一个新的 value head。用户可以通过明确恢复保存的头，在调用 train() 之前保留学习到的值。训练检查点，与 save_model 导出不同，会自动恢复 value heads。

对于使用默认线性 value head 的保存的两个 actor MAAC 设置：

from pathlib import Path
import torch
from transformers import AutoTokenizer
from comlrl.trainers.actor_critic import MAACConfig, MAACTrainer

saved_dir = Path("./saved-models/maac-run-1")
actor_dirs = [str(saved_dir / f"agent_{i}") for i in range(2)]

continued_maac = MAACTrainer(
    agents=actor_dirs,
    tokenizer=[AutoTokenizer.from_pretrained(path) for path in actor_dirs],
    critic_model=str(saved_dir / "critic"),
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAACConfig(num_agents=2, num_turns=1, num_train_epochs=3),
)
continued_maac.critics[0].value_head.load_state_dict(
    torch.load(
        saved_dir / "critic" / "value_head.pt",
        map_location="cpu",
        weights_only=True,
    )
)
continued_maac.train()
continued_maac.save_model("./saved-models/maac-run-2")

对于 IAC，与 continued_iac.agents[i].value_head 相同的 load_state_dict 步骤也适用于 use_separate_critic=False，或适用于 continued_iac.critics[i].value_head，当为 True 时。独立的 critic 和 backbones 也需要通过 critic_model 或 critics 加载。

value-head 的架构必须与保存的运行状态匹配。 IAC 需要原始的 IACConfig.value_head_hidden_dim 用于共享头部和 IACConfig.critic_value_head_hidden_dim 用于单独的头部。自定义 MAAC 头部需要原始的 model_config={"critic_value_head_hidden_dim": ...} 在 trainer 构造函数中。继续使用相同的价值学习设置也需要原始的 critic_type ("v" 或 "q")。

训练状态#

用户可以通过 train(checkpoint_dir=...) 启用训练检查点，并使用 resume_from_checkpoint 恢复与 trainer 兼容的运行状态:

# Save a checkpoint at each completed training boundary.
trainer.train(checkpoint_dir="./checkpoints", checkpoint_interval=1)

# Recreate the trainer with the same models, tokenizers, data, and callbacks.
resumed_trainer.train(
    resume_from_checkpoint="./checkpoints/checkpoint-1",
    checkpoint_dir="./checkpoints",
    checkpoint_interval=1,
)

检查点保存模型和优化器状态、训练进度、随机状态，以及，如果适用，PbRL 数据和历史。

检查点在完成的 epochs、迭代或 PbRL 阶段保存，并在完成时进行最终保存。 checkpoint_interval 控制周期性保存。恢复会重复自上次检查点以来的所有工作。

num_train_epochs 和 num_iterations 统计总运行，包括已完成的工作，并在恢复时可以增加。模型、tokenizers、数据和训练设置必须保持兼容。状态化数据集和回调可以提供 state_dict() 和 load_state_dict(); 外部服务状态不会恢复。如果 PbRL 输出目录已包含文件，则恢复的输出将进入 resume-* 子目录。

对于 W&B，可以先将保存的运行状态合并到现有的 wandb_config 中，然后再构建 trainer。这可以保留日志和输出设置，并在构建过程中避免打开一个空运行：

from comlrl.utils import read_checkpoint_metadata

checkpoint = "./checkpoints/checkpoint-1"
identity = read_checkpoint_metadata(checkpoint)["wandb"]
if identity and "id" not in wandb_config and "resume" not in wandb_config:
    wandb_config = {**wandb_config, **identity, "resume": "allow"}

# Pass wandb_config to the trainer constructor before resuming.

显式 id 或 resume 设置优先。如果检查点不包含 W&B 身份，则助手返回 None 用于 wandb，并且现有配置保持不变。

检查点默认禁用：train() 在不带检查点选项的情况下，不保存任何训练检查点。save_model() 仅导出模型；它不保存用于恢复训练所需的完整状态。