Entraînement basé sur le ratio de longueur

Ce tutoriel démontre une tâche d’entraînement simple avec deux agents, utilisant une récompense basée sur le ratio de longueur. Les deux agents reçoivent la même requête et génèrent indépendamment ; le second agent ne voit pas la réponse du premier. La récompense encourage la seconde sortie à contenir 2–3× plus de caractères que la première. Elle ne évalue pas la qualité de l’histoire ou la collaboration sémantique.

Les utilisateurs peuvent exécuter l’exemple à partir d’un dépôt installé :

python examples/story-len-ratio.py

Les exigences en mémoire dépendent du modèle, du nombre d’agents, de la taille du lot et de la longueur de la réponse. Le journalisation W&B est désactivée par défaut. Le drapeau --wandb l’active ; --wandb-project et --wandb-entity sont optionnels.

Importer les bibliothèques#

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

Préparation des données#

L’ensemble de données contient des requêtes créatives pour les deux agents.

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)

Initialisation des agents#

Le tokenizer convertit le texte en modèle tokens. Chaque agent utilise une instance de modèle distincte.

model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
agents = [AutoModelForCausalLM.from_pretrained(model_name) for _ in range(2)]

Définir la fonction de récompense#

La fonction de récompense évalue le ratio des longueurs de sortie des agents. Elle renvoie la récompense maximale (1.0) lorsque la seconde sortie contient 2–3× plus de caractères que la première. En dehors de cette plage, la récompense diminue exponentiellement avec la distance par rapport à la limite.

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

Configurer l’entraînement#

Les utilisateurs peuvent configurer l’entraînement via MAGRPOConfig, y compris le taux d’apprentissage, la taille du lot et le nombre de générations par agent par requête.

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

Créer le Trainer#

L’exemple suivant initialise MAGRPOTrainer avec les agents, la fonction de récompense et la configuration définies ci-dessus. RewardProcessors.scale multiplie les récompenses de tâche par 100. Avec la configuration par défaut advantage_normalization=True, la normalisation annule largement cet effet de mise à l’échelle positive, de sorte qu’elle ne signifie pas un gradient de politique 100 fois plus fort.

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

Exécuter l’entraînement#

Le code suivant entraîne les deux agents pour maximiser la récompense partagée et sauvegarde les modèles entraînés.

trainer.train()
trainer.save_model(f"{output_dir}/models")