Mehrstufiges Feedback

In diesem Beispiel beantworten zwei Agenten eine Aufgabe, erhalten Feedback auf ihre Antworten und beantworten dann erneut.

Einrichtung#

Die Aufgabe verlangt, dass der zweite Agent eine Antwort verfasst, die 2 bis 3 Mal so lang ist wie die Antwort des ersten Agenten.

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

Feedback zwischen den Schritten#

Nach dem ersten Schritt schreibt diese Funktion einen neuen Prompt für jeden Agenten. Der Prompt gibt die Länge der Antworten an.

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

Zwei-Schritt-Training#

num_turns=2 fügt den zweiten Schritt hinzu und external_transition verbindet die Feedback-Funktion.

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 beendet einen Zweig erst nach dem ersten Schritt, wenn die Antworten bereits das Ziel erreicht haben. Der Standardwert −0.2 würde jeden Zweig nach dem ersten Schritt beenden, da diese Belohnung niemals negativ ist.

Was passiert in jedem Schritt#

Im ersten Schritt schreibt jeder Agent 4 Antworten, die 4 gemeinsame Antworten bilden.
Die Feedback-Funktion wandelt jede gemeinsame Antwort in neue Prompts um.
Im zweiten Zug beantwortet jeder Agent erneut, sodass eine Aufgabe bis zu 16 gemeinsame Antworten hat.
Eine Antwort aus dem ersten Zug wird anhand ihrer eigenen Belohnung plus der durchschnittlichen Belohnung der nachfolgenden Antworten aus dem zweiten Zug bewertet.
Mit W&B aktiviert, zeigen turn_1/reward_mean und turn_2/reward_mean, ob die zweiten Antworten besser sind.