Inicialización del entrenamiento

El entrenamiento se inicializa a partir de un conjunto de tareas, un conjunto de tareas que comparten un formato de entrada y una recompensa. Los formatters convierten cada tarea en indicaciones para el agente, y una función de recompensa evalúa las respuestas conjuntas de los agentes.

Conjuntos de tareas#

Un conjunto de tareas es un Hugging Face Dataset o IterableDataset cuyos elementos de fila son elementos de tarea, cada uno un diccionario. El formatteur predeterminado lee el campo prompt; campos adicionales pueden contener objetivos, etiquetas, casos de prueba u otra información requerida por la función de recompensa. Trainers pasa el elemento de tarea original a la función de recompensa a través de 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"]

Aquí, ocho elementos de tarea se utilizan para el entrenamiento y dos se mantienen para la evaluación. seed asegura que la división sea repetible; no establece la semilla aleatoria para la generación del modelo. Se pueden utilizar directamente las divisiones de entrenamiento y evaluación existentes.

Los cargadores de datos de entrenamiento actuales preservan el orden de las tareas. Los usuarios pueden mezclar los elementos de tarea con train_dataset.shuffle(seed=42) antes de crear el trainer. Esto produce un orden fijo, sin reordenar cada epoch.

Indicaciones de las tareas#

Un formatteur puede compartir un formato de indicación entre agentes, o una lista de formatteurs en el orden actor puede asignar diferentes roles. Solo los campos utilizados por un formatteur entran en la indicación del modelo; la función de recompensa aún puede leer el elemento de tarea completo.

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]

Para tareas de múltiples turnos, los formatteurs también pueden aceptar external_prompts. La interfaz de transición se describe en Interacción de múltiples turnos.

Recompensas#

Tanto las recompensas basadas en reglas como los modelos de recompensa LLM se pasan a trainers a través de reward_func. La función de recompensa recibe una lista de completados por agente, en el orden de actor. Si acepta una palabra clave batch_items, trainers también pasan los elementos de la tarea. Devuelve una recompensa por respuesta conjunta. El trainers utilizado aquí lo llama para un elemento de tarea a la vez, con ese elemento disponible como batch_items[0].

reward_processor transforma cada recompensa devuelta antes del entrenamiento; por ejemplo, RewardProcessors.scale(factor=10.0) de comlrl.utils.reward_processor multiplica las recompensas por 10. Por defecto, las recompensas no se modifican.

Recompensas basadas en reglas#

Una recompensa basada en reglas calcula puntuaciones directamente a partir de las respuestas y los elementos de la tarea, por ejemplo, con comprobaciones de tareas, verificadores o marcos de prueba. La siguiente función puntúa una tarea de relación de longitud:

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

Esta función devuelve una recompensa compartida por cada par de respuestas alineado. La tarea recompensa una relación de conteo de caracteres; no mide la calidad de la escritura, y los agentes generan de forma independiente.

Los usuarios pueden comprobar el conjunto de tareas, los formateadores y la función de recompensa sin cargar los modelos:

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]

Modelos de recompensa basados en LLM#

Una función de recompensa también puede puntuar respuestas con un modelo de recompensa LLM existente; trainers utiliza solo los valores que devuelve. El siguiente esquema envuelve un modelo de recompensa con un encabezado de clasificación de un solo valor; reward_model_id es un marcador de posición:

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,

El modelo de recompensa se puede cargar una vez, fuera de la función, y colocar en un dispositivo con memoria disponible; trainers no lo mueven ni lo actualizan. Para aprender un modelo de recompensa a partir de preferencias, consulte Modelos de recompensa LLM.

Entrenamiento y Evaluación#

Los usuarios pueden proporcionar elementos de tareas de entrenamiento y evaluación por separado a través de train_dataset y eval_dataset. La evaluación utiliza los formatters de entrenamiento y la función de recompensa sin actualizar los pesos del modelo. Las salidas se muestrean con la configuración de generación configurada, por lo que las evaluaciones repetidas pueden producir diferentes puntuaciones.

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)

Para la familia MAREINFORCE, IAC y MAAC, eval_interval cuenta los elementos de tareas de entrenamiento. La evaluación se ejecuta antes del elemento cero de cada epoch, luego en el intervalo configurado; 0 habilita la evaluación manual únicamente. eval_num_samples especifica un límite positivo en los elementos de tareas evaluados, comenzando desde el principio del segmento de evaluación. Establecerlo en la longitud del segmento evalúa todos los elementos de tareas. eval_batch_size controla el tamaño del lote del cargador de datos de evaluación, no el número de respuestas generadas por el agente.

MAGRPO, IAC y MAAC generan una respuesta por agente por turno durante la evaluación. IAC y MAAC utilizan su ruta rollout, con la estructura de turno configurada, pero su método evaluate() no toma ningún argumento de conteo de muestras:

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

Sin un eval_dataset, evaluate() devuelve un diccionario vacío. Para PbRL trainers, también está disponible la evaluación manual evaluate(num_eval_samples=...); el horario de evaluación automático sigue la etapa de actualización de preferencias o actor-entrenamiento en lugar del horario de elementos de tareas anterior.

Registro#

CoMLRL admite curvas de entrenamiento en Weights & Biases, métricas específicas de la tarea y registros de evaluación locales. Los ejemplos utilizan trainer y el segmento de evaluación anterior.

Weights & Biases#

Los usuarios pueden registrar curvas de entrenamiento en Weights & Biases suministrando wandb_config al crear 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 puede autenticarse en el entorno del usuario antes de que se habilite el registro. Omitir wandb_config o establecerlo en None desactiva W&B. Los ajustes de ejecución admitidos incluyen project, entity, name (o run_name), tags y dir. El campo config_sections añade la configuración de la tarea y del experimento; config_sections.output.verbose controla la salida de progreso de trainer.

Métricas de entrenamiento#

Trainers registra las siguientes métricas de forma predeterminada, una por turno k (a partir de 1). Métricas personalizadas pueden añadir valores específicos de la tarea junto con ellas.

Con la referencia KL habilitada, reference_kl_mean y reference_kl_penalty_mean también se registran para cada turno. Las métricas de entrenamiento se registran contra env_step, el número de evaluaciones de recompensa de la respuesta conjunta, que crece más rápido que el índice de la tarea con múltiples generaciones o turnos. logging_steps establece el mínimo de la brecha env_step entre los registros de entrenamiento; las métricas de evaluación siguen eval_interval.

Métricas personalizadas#

Los usuarios pueden añadir métricas de evaluación personalizadas a MAGRPO proporcionando un eval_logger y un eval_aggregator juntos. El registrador recibe las finalizaciones en el orden [agent][task item][turn] más la original batch_items y los datos brutos prompts. El agregador devuelve métricas escalares sin el prefijo eval/; trainer añade el prefijo.

El siguiente ejemplo mide la longitud media de la respuesta en los agentes en cada turno:

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,

El diccionario resultante incluye eval/turn_1/response_chars. Estas funciones también reciben campos de tarea-elemento, permitiendo calcular métricas específicas de la tarea o grupales sin respuestas de referencia en los prompts.

IAC y MAAC en cambio aceptan metrics_callback(rollouts), que devuelve un diccionario de valores numéricos. Cada muestra rollout proporciona campos como agent_idx, prompt, completion y 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(...).

Esta función se ejecuta durante el entrenamiento y la evaluación. La evaluación promedia sus valores en función de los elementos de tarea evaluados y añade eval/; el entrenamiento registra sus valores en el momento de la colección rollout, por separado de la puerta estándar logging_steps.

Métricas locales#

La evaluación devuelve un diccionario de métricas incluso cuando W&B está desactivado. Este diccionario se puede guardar localmente para registrar una evaluación completada:

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