長さ比学習

このチュートリアルは、長さ比報酬を用いた単純な二エージェント学習タスクを示しています。両のエージェントは同じプロンプトを受け取り、独立して生成します。2番目のエージェントは、最初のエージェントの応答を見ません。この報酬は、2～3倍の文字数を最初の出力に含めるように、2番目の出力を奨励します。ストーリーの品質や意味的な協調性を評価しません。

ユーザーは、インストールされたソースチェックアウトから例を実行できます:

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番目の出力が最初の出力の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")