学習の初期化

学習は、入力形式と報酬を共有するタスクのセットから開始されます。フォーマッタは、各タスクをエージェントのプロンプトに変換し、報酬関数はエージェントの共同応答を評価します。

タスクセット#

タスクセットは、 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"]

ここでは、8つのタスクアイテムが学習に使用され、2つが評価に使用されます。seedは、分割を再現可能にします。モデル生成のランダムシードを設定しません。既存の学習および評価の分割を直接使用できます。

現在の学習データローダーは、タスクの順序を保持します。ユーザーは、trainerを作成する前に、train_dataset.shuffle(seed=42)を使用してタスクアイテムをシャッフルできます。これにより、各epochに対して固定された順序が生成されます。

タスクからのプロンプト#

1つのフォーマッタは、複数のエージェント間でプロンプト形式を共有できます。または、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 に渡されます。報酬関数は、エージェントごとに 1 つの完了リストを受け取り、actor の順序で受け取ります。 batch_items キーワードを受け入れる場合、trainers はタスクアイテムも渡します。 1 つの共同応答ごとに 1 つの報酬を返します。 trainers は、この報酬関数を 1 つのタスクアイテムごとに呼び出し、そのアイテムが 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

この関数は、対応する応答の 1 つ共有された報酬を返します。タスクは文字数を比率で評価します。これは、執筆の品質を測定するものではなく、エージェントは独立して生成されます。

ユーザーは、モデルをロードせずに、タスクセット、フォーマッタ、および報酬関数を確認できます。

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,

報酬モデルは、関数外で 1 回ロードし、メモリの空いているデバイスに配置できます。 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 は、評価中にエージェントごとにターンごとに 1 つの応答を生成します。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.

カスタム指標#

ユーザーは、MAGRPO にカスタム評価指標を追加するために、eval_logger と eval_aggregator を一緒に提供できます。ロガーは、[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/ を追加します。トレーニングは、標準の logging_steps ゲートとは別に、rollout 収集時に値を記録します。

ローカル指標#

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")