다중 턴 피드백

이 예제에서, 두 에이전트는 작업을 수행하고, 답변에 대한 피드백을 받고, 다시 답변합니다.

설정#

작업은 두 번째 에이전트에게 첫 번째 에이전트의 답변보다 2에서 3배 더 긴 답변을 작성하도록 요청합니다.

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

턴 간 피드백#

첫 번째 턴 이후, 이 함수는 각 에이전트에게 새로운 프롬프트를 작성합니다. 프롬프트는 답변의 길이를 보고합니다.

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

2 턴 훈련#

num_turns=2 는 두 번째 턴을 추가하고, external_transition 는 피드백 함수를 연결합니다.

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 는 첫 번째 턴 이후에만 특정 분기를 종료합니다. 즉, 답변이 이미 목표에 도달한 경우에만, 기본값인 −0.2는 첫 번째 턴 이후 모든 분기를 종료합니다. 왜냐하면 이 보상은 항상 음수가 아니기 때문입니다.

각 단계에서 발생하는 일#

첫 번째 턴에서, 각 에이전트는 4개의 답변을 작성하여 4개의 공동 답변을 형성합니다.
피드백 함수는 각 공동 답변을 새로운 프롬프트로 변환합니다.
두 번째 턴에서, 각 에이전트는 다시 답변을 제공하므로, 하나의 작업은 최대 16개의 공동 응답을 가질 수 있습니다.
첫 번째 턴의 답변은 자체 보상과 그 뒤에 오는 두 번째 턴의 답변의 평균 보상으로 평가됩니다.
W&B 활성화 시, turn_1/reward_mean 와 turn_2/reward_mean 는 두 번째 답변이 더 나은지 여부를 보여줍니다.