长度比训练

本教程演示了一个简单的双智能体训练任务，使用长度比奖励。两个智能体都接收相同的提示并独立生成；第二个智能体看不到第一个智能体的响应。奖励鼓励第二个输出包含 2–3 倍第一个输出的字符数。它不评估故事质量或语义协作。

用户可以在已安装的源代码检出中运行示例：

python examples/story-len-ratio.py

内存要求取决于模型、智能体数量、批大小和响应长度。默认情况下，W&B 记录已禁用。--wandb 标志启用它；--wandb-project 和 --wandb-entity 是可选的。

导入库#

import math
from functools import partial
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from comlrl.utils.reward_processor import RewardProcessors
from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer

数据集准备#

数据集包含两个智能体的创意提示。

train_data = {
    "prompt": [
        "Describe a city in the clouds:",
        "Invent a new holiday and explain it:",
        "Write a bedtime story for a dragon:",
        "Explain how teleportation might work:",
        "Tell a joke 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:",
    ]
}
train_dataset = Dataset.from_dict(train_data)

智能体初始化#

tokenizer 将文本转换为模型 tokens。每个智能体使用单独的模型实例。

model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
agents = [AutoModelForCausalLM.from_pretrained(model_name) for _ in range(2)]

定义奖励函数#

奖励函数评估智能体的输出长度比。当第二个输出包含 2–3 倍第一个输出的字符数时，它返回最大奖励 (1.0)。超出此范围，奖励以距离最近边界的指数衰减。

def proper_length_ratio_reward(
    completions1, completions2, target_min=2.0, target_max=3.0
):
    rewards = []
    for c1, c2 in zip(completions1, completions2):
        len1, len2 = len(c1), len(c2)

        if len1 == 0:
            rewards.append(0.0)
            continue

        ratio = len2 / len1

        if target_min <= ratio <= target_max:
            reward = 1.0
        else:
            if ratio < target_min:
                distance = target_min - ratio
            else:
                distance = ratio - target_max

            reward = math.exp(-distance)

        rewards.append(float(reward))

    return rewards

配置训练#

用户可以通过 MAGRPOConfig 配置训练，包括学习率、批次大小和每个代理每提示的生成次数。

output_dir = "./magrpo_multi_reward_output"
config = MAGRPOConfig(
    num_train_epochs=3,
    agent_learning_rate=5e-5,
    logging_steps=10,
    num_generations=8,
    num_turns=1,
    max_new_tokens=128,
)

创建 Trainer#

以下示例初始化 MAGRPOTrainer，其中包含上述定义的代理、奖励函数和配置。 RewardProcessors.scale 将任务奖励乘以 100。默认情况下，advantage_normalization=True 会很大程度上抵消这种正向缩放，因此它并不意味着 100 倍更强的策略梯度。

configured_reward_func = partial(
    proper_length_ratio_reward, target_min=2, target_max=3
)

trainer = MAGRPOTrainer(
    agents=agents,
    reward_func=configured_reward_func,
    reward_processor=RewardProcessors.scale(factor=100.0),
    args=config,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)

运行训练#

以下代码训练两个代理以最大化共享奖励并保存训练后的模型。

trainer.train()
trainer.save_model(f"{output_dir}/models")