Entrenamiento basado en la relación de longitud

Este tutorial demuestra una tarea de entrenamiento simple con dos agentes, utilizando una recompensa basada en la relación de longitud. Ambos agentes reciben la misma instrucción y generan de forma independiente; el segundo agente no ve la respuesta del primer agente. La recompensa incentiva que la segunda salida contenga 2–3× más caracteres que la primera. No evalúa la calidad de la historia ni la colaboración semántica.

Los usuarios pueden ejecutar el ejemplo desde una copia de código instalada:

python examples/story-len-ratio.py

Los requisitos de memoria dependen del modelo, el número de agentes, el tamaño del lote y la longitud de la respuesta. El registro de W&B está deshabilitado por defecto. La bandera --wandb lo habilita; --wandb-project y --wandb-entity son opcionales.

Importar bibliotecas#

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

Preparación del conjunto de datos#

El conjunto de datos contiene instrucciones creativas para ambos agentes.

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)

Inicialización del agente#

El tokenizer convierte el texto en el modelo tokens. Cada agente utiliza una instancia del modelo separada.

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

Definir la función de recompensa#

La función de recompensa evalúa la relación de longitud de las salidas de los agentes. Devuelve la recompensa máxima (1.0) cuando la segunda salida contiene 2–3× más caracteres que la primera. Fuera de este rango, la recompensa decae exponencialmente con la distancia al límite más cercano.

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

Configurar el entrenamiento#

Los usuarios pueden configurar el entrenamiento a través de MAGRPOConfig, incluyendo la tasa de aprendizaje, el tamaño del lote y el número de generaciones por agente por prompt.

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

Crear el Trainer#

El siguiente ejemplo inicializa MAGRPOTrainer con los agentes, la función de recompensa y la configuración definidas anteriormente. RewardProcessors.scale multiplica las recompensas de la tarea por 100. Con el advantage_normalization=True predeterminado, la normalización cancela en gran medida esta escala positiva, por lo que no implica un gradiente de política 100 veces más fuerte.

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

Ejecutar el entrenamiento#

El siguiente código entrena a ambos agentes para maximizar la recompensa compartida y guarda los modelos entrenados.

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