Retroalimentación en Múltiples Turnos

En este ejemplo, dos agentes responden a una tarea, reciben retroalimentación sobre sus respuestas y luego responden de nuevo.

Configuración#

La tarea pide al segundo agente que escriba una respuesta de 2 a 3 veces más larga que la respuesta del primer agente.

import math

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer

model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
train_dataset = Dataset.from_dict({
    "prompt": [
        "Describe a city in the clouds:",
        "Invent a new holiday and explain it:",
        "Write a bedtime story for a dragon:",
        "Design a playground on Mars:",
    ]
})


def length_ratio_reward(completions1, completions2):
    rewards = []
    for first, second in zip(completions1, completions2):
        ratio = len(second) / max(len(first), 1)
        distance = 0.0 if 2 <= ratio <= 3 else min(abs(ratio - 2), abs(ratio - 3))
        rewards.append(math.exp(-distance))
    return rewards

Retroalimentación entre Turnos#

Después del primer turno, esta función escribe una nueva instrucción para cada agente. La instrucción informa sobre la longitud de las respuestas.

def length_feedback(prompt, agent_completions, num_agents, **history):
    lengths = [len(text) for text in agent_completions]
    ratio = lengths[1] / max(lengths[0], 1)
    return [
        f"{prompt}\n\nYour last answer had {lengths[i]} characters. "
        f"The second answer was {ratio:.1f} times as long as the first; "
        "the goal is 2 to 3 times. Please write an improved answer."
        for i in range(num_agents)
    ]

Entrenamiento de Dos Turnos#

num_turns=2 añade el segundo turno, y external_transition conecta la función de retroalimentación.

config = MAGRPOConfig(
    num_turns=2,
    num_train_epochs=2,
    num_generations=4,
    max_new_tokens=64,
    logging_steps=1,
    early_termination_threshold=0.9,
)
trainer = MAGRPOTrainer(
    agents=[AutoModelForCausalLM.from_pretrained(model_name) for _ in range(2)],
    tokenizer=tokenizer,
    reward_func=length_ratio_reward,
    external_transition=length_feedback,
    train_dataset=train_dataset,
    args=config,
)
trainer.train()

early_termination_threshold=0.9 finaliza una rama después del primer turno solo cuando sus respuestas ya alcanzan el objetivo. El valor predeterminado, −0.2, finalizaría cada rama después del primer turno, ya que esta recompensa nunca es negativa.

Qué ocurre en cada paso#

En el primer turno, cada agente escribe 4 respuestas, que forman 4 respuestas conjuntas.
La función de retroalimentación convierte cada respuesta conjunta en nuevas instrucciones.
En la segunda ronda, cada agente responde nuevamente, por lo que una tarea puede tener hasta 16 respuestas conjuntas.
Una respuesta de la primera ronda se califica según su propia recompensa más la recompensa promedio de las respuestas de la segunda ronda que la siguen.
Con W&B habilitado, turn_1/reward_mean y turn_2/reward_mean muestran si las respuestas de la segunda ronda son mejores.