길이-비율 훈련

이 튜토리얼은 길이-비율 보상을 사용하는 간단한 두 에이전트 훈련 작업을 보여줍니다. 두 에이전트는 동일한 프롬프트를 받고 독립적으로 생성합니다; 두 번째 에이전트는 첫 번째 에이전트의 응답을 보지 않습니다. 이 보상은 두 번째 출력이 첫 번째 출력보다 2–3배 더 많은 문자를 포함하도록 장려합니다. 이는 스토리 품질이나 의미론적 협업을 평가하지 않습니다.

사용자는 설치된 소스 체크아웃에서 예제를 실행할 수 있습니다:

python examples/story-len-ratio.py

메모리 요구 사항은 모델, 에이전트 수, 배치 크기, 응답 길이에 따라 달라집니다. W&B 로깅은 기본적으로 비활성화되어 있습니다. --wandb 플래그를 활성화하면 활성화됩니다; --wandb-project 및 --wandb-entity는 선택 사항입니다.

라이브러리 가져오기#

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

데이터셋 준비#

데이터셋에는 두 에이전트 모두를 위한 창의적인 프롬프트가 포함되어 있습니다.

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)

에이전트 초기화#

tokenizer는 텍스트를 모델 tokens로 변환합니다. 각 에이전트는 별도의 모델 인스턴스를 사용합니다.

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

보상 함수 정의#

보상 함수는 에이전트의 출력 길이 비율을 평가합니다. 두 번째 출력이 첫 번째 출력보다 2–3배 더 많은 문자를 포함하는 경우 최대 보상(1.0)을 반환합니다. 이 범위 밖에서는 보상이 가장 가까운 경계에서 거리에 따라 지수적으로 감소합니다.

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

훈련 구성#

사용자는 MAGRPOConfig 를 통해 훈련을 구성할 수 있으며, 여기에는 학습률, 배치 크기, 그리고 각 에이전트당 프롬프트당 생성 횟수가 포함됩니다.

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

Trainer 생성#

다음 예시는 위에서 정의된 에이전트, 보상 함수, 그리고 구성으로 MAGRPOTrainer 를 초기화합니다. RewardProcessors.scale 는 작업 보상을 100 배 곱합니다. 기본 advantage_normalization=True 의 경우, 정규화가 이러한 양의 스케일링을 크게 상쇄하므로, 이는 100 배 더 강력한 정책 경사라는 의미가 아닙니다.

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

훈련 실행#

다음 코드는 두 에이전트를 공유 보상을 최대화하도록 훈련하고, 훈련된 모델을 저장합니다.

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