Initialisierung des Trainings

Das Training wird von einem Aufgaben-Set initialisiert, einer Klasse von Aufgaben, die ein einheitliches Eingabeformat und eine Belohnung teilen. Formatierer wandeln jede Aufgabe in Agent-Prompts um, und eine Belohnungsfunktion bewertet die gemeinsamen Antworten der Agenten.

Aufgaben-Sets#

Ein Aufgaben-Set ist ein Hugging Face Dataset oder IterableDataset, dessen Zeilen Aufgaben sind, wobei jede Aufgabe ein Dictionary ist. Der Standard-Formatierer liest das Feld prompt; zusätzliche Felder können Ziele, Labels, Testfälle oder andere Informationen enthalten, die von der Belohnungsfunktion benötigt werden. Trainers übergibt das ursprüngliche Aufgaben-Item über batch_items an die Belohnungsfunktion.

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

Hier werden acht Aufgaben-Items für das Training und zwei für die Bewertung verwendet. seed stellt die Wiederholbarkeit sicher; es setzt keine Zufalls-Seed für die Modellgenerierung. Bestehende Trainings- und Bewertungs-Splits können direkt verwendet werden.

Die aktuellen Trainings-Data-Loader behalten die Aufgaben-Reihenfolge bei. Benutzer können die Aufgaben-Items mit train_dataset.shuffle(seed=42) vor dem Erstellen des trainer mischen. Dies erzeugt eine feste Reihenfolge, ohne jedes epoch neu zu mischen.

Prompts aus Aufgaben#

Ein Formatierer kann ein Prompt-Format für mehrere Agenten teilen, oder eine Liste von Formatierern in der Reihenfolge actor kann unterschiedliche Rollen zuweisen. Nur die Felder, die von einem Formatierer verwendet werden, gelangen in den Modell-Prompt; die Belohnungsfunktion kann das vollständige Aufgaben-Item weiterhin lesen.

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]

Für mehrstufige Aufgaben können Formatierer auch external_prompts akzeptieren. Die Übergangs-Schnittstelle wird in Mehrstufige Interaktion beschrieben.

Belohnungen#

Sowohl regelbasierte Belohnungen als auch LLM-Belohnungsmodelle werden über reward_func an trainers übergeben. Die Belohnungsfunktion erhält eine Liste von Vervollständigungen pro Agent, in der Reihenfolge von actor. Wenn sie ein batch_items-Schlüsselwort akzeptiert, übermitteln trainers auch die Aufgaben. Sie gibt eine Belohnung pro gemeinsamer Antwort zurück. trainers verwendet dies für ein Aufgaben-Element gleichzeitig, wobei dieses verfügbar ist als batch_items[0].

reward_processor transformiert jede zurückgegebene Belohnung vor dem Training; zum Beispiel multipliziert RewardProcessors.scale(factor=10.0) von comlrl.utils.reward_processor Belohnungen mit 10. Standardmäßig bleiben Belohnungen unverändert.

Regelbasierte Belohnungen#

Eine regelbasierte Belohnung berechnet direkt Werte aus den Antworten und Aufgaben, z. B. durch Aufgabenprüfungen, Verifizierer oder Testumgebungen. Die folgende Funktion bewertet eine Aufgaben-Längenverhältnis-Aufgabe:

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

Diese Funktion gibt eine gemeinsame Belohnung für jedes ausgerichtete Paar von Antworten zurück. Die Aufgaben belohnen ein Zeichen-Zahlen-Verhältnis; sie messen nicht die Schreibqualität, und die Agenten generieren unabhängig voneinander.

Benutzer können den Aufgaben-Set, die Formatierer und die Belohnungsfunktion ohne das Laden von Modellen überprüfen:

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]

Belohnungsmodelle basierend auf LLM#

Eine Belohnungsfunktion kann auch Antworten mit einem bestehenden LLM-Belohnungsmodell bewerten; trainers verwenden nur die Werte, die es zurückgibt. Der folgende Entwurf umschließt ein Belohnungsmodell mit einem Klassifikations-Head mit einem einzelnen Wert; reward_model_id ist ein Platzhalter:

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,

Das Belohnungsmodell kann einmal geladen werden, außerhalb der Funktion, und auf einem Gerät mit freiem Speicher platziert werden; trainers verschieben oder aktualisieren es nicht. Um ein Belohnungsmodell aus Präferenzen zu lernen, siehe Belohnungsmodelle basierend auf LLM.

Training und Bewertung#

Benutzer können Trainings- und Evaluationsaufgaben separat über train_dataset und eval_dataset bereitstellen. Die Evaluation verwendet die Trainingsformatierer und die Belohnungsfunktion, ohne die Modellgewichte zu aktualisieren. Die Ausgaben werden mit den konfigurierten Generierungseinstellungen abgetastet, sodass wiederholte Bewertungen unterschiedliche Ergebnisse liefern können.

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)

Für die MAREINFORCE-Familie, IAC und MAAC zählt eval_interval Trainingsaufgaben. Die Evaluation wird vor Aufgabe 0 jeder epoch durchgeführt, dann in dem konfigurierten Intervall; 0 ermöglicht nur manuelle Evaluation. eval_num_samples legt eine positive Grenze für evaluierte Aufgaben fest, beginnend mit dem Anfang des Evaluationssplits. Das Setzen auf die Länge des Splits bewertet alle Aufgaben. eval_batch_size steuert die Größe des Datenladers für die Evaluation, nicht die Anzahl der generierten Antworten pro Agent.

MAGRPO, IAC und MAAC generieren eine Antwort pro Agent pro Runde während der Evaluation. IAC und MAAC verwenden ihren rollout-Pfad, mit der konfigurierten Rundenstruktur, aber ihre evaluate()-Methode benötigt keinen Argument für die Stichprobengröße:

# For an existing IACTrainer or MAACTrainer:
ac_trainer.args.eval_num_samples = len(eval_dataset)
metrics = ac_trainer.evaluate()

Ohne eval_dataset, evaluate() gibt ein leeres Dictionary zurück. Für PbRL trainers ist auch eine manuelle evaluate(num_eval_samples=...) verfügbar; der automatische Evaluationsplan folgt der Präferenz-Update- oder actor-Trainingsphase, anstatt dem oben genannten Aufgabenplan.

Protokollierung#

CoMLRL unterstützt Trainingskurven in Weights & Biases, aufgabenbezogene Metriken und lokale Evaluationsaufzeichnungen. Die Beispiele verwenden trainer und den oben genannten Evaluationssplit.

Weights & Biases#

Benutzer können Trainingskurven in Weights & Biases aufzeichnen, indem sie wandb_config bereitstellen, wenn sie trainer erstellen:

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 kann in der Umgebung des Benutzers authentifiziert werden, bevor das Protokollieren aktiviert wird. Das Auslassen von wandb_config oder das Setzen auf None deaktiviert W&B. Unterstützte Ausführungssettings sind project, entity, name (oder run_name), tags und dir. Das Feld config_sections fügt Task- und Experiment-Einstellungen zur aufgezeichneten Konfiguration hinzu; config_sections.output.verbose steuert die trainer-Fortschrittsausgabe.

Trainingsmetriken#

Trainers protokolliert standardmäßig die folgenden Metriken, einmal pro Runde k (beginnend bei 1). Benutzerdefinierte Metriken können zusätzliche Task-spezifische Werte hinzufügen.

Mit aktiviertem Referenz-KL werden außerdem reference_kl_mean und reference_kl_penalty_mean für jede Runde protokolliert. Trainingsmetriken werden gegen env_step, die Anzahl der gemeinsamen-Aktions-Belohnungsbewertungen, protokolliert, die mit mehreren Generationen oder Runden schneller wächst als der Task-Index. logging_steps legt die minimale env_step-Lücke zwischen Trainingsprotokollen fest; Evaluationsmetriken folgen eval_interval.

Benutzerdefinierte Metriken#

Benutzer können benutzerdefinierte Evaluationsmetriken zu MAGRPO hinzufügen, indem sie einen eval_logger und einen eval_aggregator zusammen liefern. Der Logger empfängt Ergebnisse in der Reihenfolge [agent][task item][turn] sowie die ursprünglichen batch_items und rohen prompts. Der Aggregator gibt skalare Metriken ohne das Präfix eval/ zurück; trainer fügt es hinzu.

Das folgende Beispiel misst die durchschnittliche Antwortlänge über Agenten bei jeder Runde:

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,

Das resultierende Dictionary enthält eval/turn_1/response_chars. Diese Callbacks erhalten außerdem Task-Item-Felder, wodurch Task-spezifische Genauigkeits- oder Gruppierungsmetriken ohne Referenzantworten in den Prompts berechnet werden können.

IAC und MAAC akzeptieren stattdessen metrics_callback(rollouts), das ein Dictionary mit numerischen Werten zurückgibt. Jedes rollout-Beispiel bietet Felder wie agent_idx, prompt, completion und 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(...).

Dieser Callback wird sowohl während des Trainings als auch der Evaluierung ausgeführt. Die Evaluierung mittelt ihre Werte über die ausgewerteten Task-Items und fügt eval/ hinzu; das Training protokolliert seine Werte zur rollout-Zeit, getrennt von der Standard-logging_steps-Gate.

Lokale Metriken#

Die Evaluierung gibt ein Metriken-Dictionary zurück, auch wenn W&B deaktiviert ist. Dieses Dictionary kann lokal gespeichert werden, um eine abgeschlossene Evaluierung aufzuzeichnen:

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