Initialisation de l'entraînement

L’entraînement est initialisé à partir d’un ensemble de tâches, une classe de tâches qui partagent un format d’entrée et une récompense. Les formatteurs transforment chaque tâche en invites pour l’agent, et une fonction de récompense évalue les réponses conjointes des agents.

Ensembles de tâches#

Un ensemble de tâches est un Hugging Face Dataset ou IterableDataset dont les lignes sont des éléments de tâche, chacun un dictionnaire. Le formatteur par défaut lit le champ prompt; des champs supplémentaires peuvent contenir des cibles, des étiquettes, des cas de test ou d’autres informations requises par la fonction de récompense. Trainers transmet l’élément de tâche original à la fonction de récompense via 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"]

Ici, huit éléments de tâche sont utilisés pour l’entraînement et deux sont conservés pour l’évaluation. seed assure la reproductibilité; il ne définit pas la graine aléatoire pour la génération du modèle. Les divisions d’entraînement et d’évaluation existantes peuvent être utilisées directement.

Les chargeurs de données d’entraînement actuels conservent l’ordre des tâches. Les utilisateurs peuvent mélanger les éléments de tâche avec train_dataset.shuffle(seed=42) avant de créer le trainer. Cela produit un ordre fixe, sans réorganiser chaque epoch.

Prompts des tâches#

Un formatteur peut partager un format de prompt entre plusieurs agents, ou une liste de formatteurs dans l’ordre actor peut attribuer différents rôles. Seuls les champs utilisés par un formatteur entrent dans le prompt du modèle; la fonction de récompense peut toujours lire l’élément de tâche complet.

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]

Pour les tâches multi-tours, les formatteurs peuvent également accepter external_prompts. L’interface de transition est décrite dans Interaction multi-tours.

Récompenses#

Les récompenses basées sur des règles et les modèles de récompense LLM sont transmis à trainers via reward_func. La fonction de récompense reçoit une liste de complétions par agent, dans l’ordre actor. Si elle accepte un mot-clé batch_items, trainers transmet également les éléments de tâche. Elle renvoie une seule récompense par réponse conjointe. Le trainers l’utilise pour une tâche à la fois, avec cet élément disponible comme batch_items[0].

reward_processor transforme chaque récompense avant l’entraînement ; par exemple, RewardProcessors.scale(factor=10.0) multiplie les récompenses par 10, provenant de comlrl.utils.reward_processor. Par défaut, les récompenses ne sont pas modifiées.

Récompenses basées sur des règles#

Une récompense basée sur des règles calcule directement les scores à partir des réponses et des éléments de tâche, par exemple avec des vérifications de tâche, des vérificateurs ou des ensembles de tests. La fonction suivante attribue une note à une tâche de rapport :

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

Cette fonction renvoie une seule récompense partagée pour chaque paire de réponses alignée. La tâche récompense un rapport de caractères ; elle ne mesure pas la qualité de l’écriture, et les agents génèrent indépendamment.

Les utilisateurs peuvent vérifier l’ensemble de tâches, les formatteurs et la fonction de récompense sans charger de modèles :

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]

Modèles de récompense basés sur LLM#

Une fonction de récompense peut également attribuer des scores aux réponses à l’aide d’un modèle de récompense existant LLM ; trainers utilise uniquement les valeurs qu’il renvoie. Le schéma suivant enveloppe un modèle de récompense avec un tête de classification à un seul score ; reward_model_id est un espace réservé :

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,

Le modèle de récompense peut être chargé une fois, en dehors de la fonction, et placé sur un appareil avec suffisamment de mémoire ; trainers ne le déplace pas ni ne le met à jour. Pour apprendre un modèle de récompense à partir de préférences, voir Modèles de récompense LLM.

Entraînement et évaluation#

Les utilisateurs peuvent fournir des éléments de tâche d’entraînement et d’évaluation séparément via train_dataset et eval_dataset. L’évaluation utilise les formatteurs de tâche d’entraînement et la fonction de récompense sans mettre à jour les poids du modèle. Les sorties sont échantillonnées avec les paramètres de génération configurés, de sorte que des évaluations répétées peuvent produire des scores différents.

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)

Pour la famille MAREINFORCE, IAC et MAAC, eval_interval compte les éléments de tâche d’entraînement. L’évaluation s’exécute avant l’élément zéro de chaque epoch, puis à l’intervalle configuré ; 0 permet uniquement une évaluation manuelle. eval_num_samples spécifie une limite positive sur les éléments de tâche évalués, en commençant au début de la division d’évaluation. Définir cette valeur sur la longueur de la division évalue tous les éléments de tâche. eval_batch_size contrôle la taille du lot du chargeur de données d’évaluation, et non le nombre de réponses générées par l’agent.

MAGRPO, IAC et MAAC génèrent une réponse par agent par tour pendant l’évaluation. IAC et MAAC utilisent leur chemin rollout, avec la structure de tour configurée, mais leur méthode evaluate() ne prend aucun argument de nombre d’échantillons :

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

Sans un eval_dataset, evaluate() renvoie un dictionnaire vide. Pour PbRL trainers, une évaluation manuelle evaluate(num_eval_samples=...) est également disponible ; le calendrier d’évaluation automatique suit la phase de mise à jour des préférences ou actor-entraînement plutôt que le calendrier des éléments de tâche ci-dessus.

Journalisation#

CoMLRL prend en charge les courbes d’entraînement dans Weights & Biases, les métriques spécifiques aux tâches et les enregistrements d’évaluation locaux. Les exemples utilisent trainer et la division d’évaluation ci-dessus.

Weights & Biases#

Les utilisateurs peuvent enregistrer les courbes d’entraînement dans Weights & Biases en fournissant wandb_config lors de la création de trainer :

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 peut être authentifié dans l’environnement de l’utilisateur avant que l’enregistrement ne soit activé. L’omission de wandb_config ou sa définition sur None désactive W&B. Les paramètres de l’exécution pris en charge incluent project, entity, name (ou run_name), tags et dir. Le champ config_sections ajoute les paramètres de tâche et d’expérience à la configuration enregistrée ; config_sections.output.verbose contrôle la sortie de progression de trainer.

Métriques d’entraînement#

Trainers enregistre par défaut les métriques suivantes, une fois par tour k (à partir de 1) : Métriques personnalisées peuvent y ajouter des valeurs spécifiques à la tâche.

En référence KL activé, reference_kl_mean et reference_kl_penalty_mean sont également enregistrés pour chaque tour. Les métriques d’entraînement sont enregistrées par rapport env_step, au nombre d’évaluations de récompense d’action conjointe, qui augmente plus rapidement que l’index de la tâche avec plusieurs générations ou tours. logging_steps définit l’écart minimum env_step entre les logs d’entraînement ; les métriques d’évaluation suivent eval_interval.

Métriques personnalisées#

Les utilisateurs peuvent ajouter des métriques d’évaluation personnalisées à MAGRPO en fournissant un eval_logger et un eval_aggregator ensemble. Le logger reçoit les complétions dans l’ordre [agent][task item][turn], ainsi que la valeur originale batch_items et les données brutes prompts. L’agrégateur renvoie des métriques scalaires sans le préfixe eval/ ; trainer ajoute ce préfixe.

L’exemple suivant mesure la longueur moyenne de la réponse des agents à chaque tour :

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,

Le dictionnaire résultant inclut eval/turn_1/response_chars. Ces callbacks reçoivent également des champs de tâche, permettant de calculer des métriques spécifiques à la tâche ou des métriques de groupe sans avoir besoin de réponses de référence dans les invites.

IAC et MAAC acceptent à la place metrics_callback(rollouts), qui renvoie un dictionnaire de valeurs numériques. Chaque échantillon rollout fournit des champs tels que agent_idx, prompt, completion et 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(...).

Cette callback s’exécute à la fois pendant l’entraînement et l’évaluation. L’évaluation calcule ses valeurs en moyenne sur les éléments de tâche évalués et ajoute eval/ ; l’entraînement enregistre ses valeurs au moment de la collecte rollout, séparément de la porte standard logging_steps.

Métriques locales#

L’évaluation renvoie un dictionnaire de métriques même lorsque W&B est désactivé. Ce dictionnaire peut être enregistré localement pour enregistrer une évaluation terminée:

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