学習対象 | CoMLRL

CoMLRL は LLM エージェントの学習をサポートし、LLM をベースとした critics または報酬モデルを使用した学習を容易にします。エージェントと critics は、Hugging Face またはローカルディレクトリからの同種または異種モデルとして設定できます。ユーザーは、同種設定の場合は Hugging Face モデル識別子を agent_model/critic_model を通じて指定し、異種設定の場合は agents/critics のリストを提供できます。学習済みの対象は、さらなる学習のために新しい trainer に保存およびロードできます。

LLM エージェント#

同種エージェント#

同種のエージェントは、agent_model を単一のモデル識別子に設定し、agents=None を残すことで初期化されます。 trainers で、別個の num_agents コンストラクタ引数を使用する場合、その引数は対応する設定値と共に設定できます。設定のみがそれを上書きしません。

次の例は、3 つの 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 にモデル識別子またはロードされたモデルオブジェクトを提供することで、異種エージェントを初期化できます。tokenizers は、tokenizer を通じて提供されます。 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 は、critic_model または 1 つの要素の critics リストを通じて読み込むことができます。これはサポートされている因果言語モデルである必要があります。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 を使用すると、ユーザーはエージェントごとに別のコピーを読み込むために、1 つの critic_model 識別子または critics の長さの num_agents のリストを提供できます。各個別の critic は、token IDs を actor の tokenizer から受け取り、互換性がある必要があります。 use_separate_critic=False を使用すると、各 actor と critic は、LLM backbone を共有し、一緒に読み込まれます。この場合、critic_model と critics は未設定 (null または None) です。エージェントと同様に、両方の critic_model と critics を提供するには、モデル名と長さが一致する必要があります。そうでない場合は、トレーニング中にエラーが発生します。

以下の例では、同じモデルに基づいて、Qwen/Qwen2.5-Coder-3B と Qwen/Qwen2.5-Coder-7B actors を個別の critics で読み込みます:

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 は agent_model に対して tokenizer を自動的にロードしますが、agents が提供されている場合は、明示的な tokenizer が必要です。ロードされたモデルオブジェクトには、tokenizer または tokenizer のリストが、actors と互換性がある必要があります。異種 actors は、それぞれ独自の互換性のある tokenizer を使用する場合、異なる語彙を持つ可能性があります。別の critics は、actor の token IDs (actor の最初の IDs は MAAC 用) を受け入れる必要があります。

LLM 報酬モデル#

In 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 のみを学習します.

次の例は、2 つの Qwen/Qwen2.5-Coder-3B エージェントに対して Qwen/Qwen2.5-Coder-7B 報酬モデルをロードする方法を示しています:

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

保存されたモデルからのトレーニング#

保存された 2 つの actor および MAGRPO 設定からトレーニングを継続するには、元の役割順序にある actors ディレクトリを供給します。親 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 は、新しいトレーニング実行における 3 つの epochs を意味します。データセット、報酬関数、トレーニング構成、およびカスタムフォーマット、報酬プロセッサ、または外部遷移は、再度提供する必要があります。 save_model はこれらを保存しません。

他の trainers も、tokenizer および actor の要件を満たす範囲で、同じローカルディレクトリのアプローチをサポートします。 Preference trainers は、新しい実行構成に基づいて選好の収集と報酬学習を開始します。

IAC および MAAC は、value heads を個別に value_head.pt として保存します。コンストラクタを通じて critic または共有された actor ディレクトリをロードすると、言語モデル backbone が復元され、新しい value head が作成されます。ユーザーは、train() を呼び出す前に、保存されたヘッドを明示的に復元することで、学習された値を保持できます。訓練チェックポイントは、save_model のエクスポートとは異なり、value heads を自動的に復元します。

デフォルトの線形 value head を使用した、保存された 2 つの 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, 同じ load_state_dict 適用される continued_iac.agents[i].value_head または 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 ヘッドは、trainer コンストラクタで元の model_config={"critic_value_head_hidden_dim": ...} を必要とします。同様の価値学習の設定を継続する場合も、元の 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 の場合、保存された実行の識別子は、trainer が構築される前に既存の wandb_config にマージできます。これにより、ログと出力設定が保持され、構築中に空の実行が開かれるのを回避できます:

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() はモデルのみのエクスポートを行い、トレーニングを再開するために必要な完全な状態は保存しません。