학습 초기화

학습은 작업 세트에서 초기화됩니다. 작업 세트는 입력 형식과 보상을 공유하는 작업의 클래스입니다. 포맷터는 각 작업을 에이전트 프롬프트로 변환하고, 보상 함수는 에이전트의 공동 응답을 평가합니다.

작업 세트#

작업 세트는 Hugging Face, Dataset 또는 IterableDataset이며, 행은 작업 항목으로, 각 항목은 딕셔너리입니다. 기본 포맷터는 prompt 필드를 읽습니다. 추가 필드는 보상 함수에 필요한 대상, 레이블, 테스트 케이스 또는 기타 정보를 포함할 수 있습니다. Trainers는 원래 작업 항목을 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"]

여기서 8개의 작업 항목이 학습에 사용되고 2개가 평가에 사용됩니다. seed는 분할을 반복 가능하게 만듭니다. 모델 생성에 대한 무작위 시드를 설정하지 않습니다. 기존 학습 및 평가 분할을 직접 사용할 수 있습니다.

현재 훈련 데이터 로더는 작업 순서를 유지합니다. 사용자는 train_dataset.shuffle(seed=42) 를 사용하여 작업 항목을 섞은 후 trainer 를 생성할 수 있습니다. 이를 통해 고정된 순서를 생성하며, 각 epoch 를 다시 섞지 않습니다.

작업에서 프롬프트#

하나의 포맷터는 에이전트 간에 프롬프트 형식을 공유할 수 있습니다. 또는 actor에서 순서대로 지정된 포맷터 목록은 다른 역할을 할당할 수 있습니다. 포맷터에 사용되는 필드만 모델 프롬프트에 들어갑니다. 보상 함수는 여전히 전체 작업 항목을 읽을 수 있습니다.

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]

다중 턴 작업의 경우, 포맷터는 또한 external_prompts를 수락할 수 있습니다. 전환 인터페이스는 다중 턴 상호 작용에서 설명되어 있습니다.

보상#

규칙 기반 보상과 LLM 보상 모델은 다음으로 전달됩니다 trainers 을 통해 reward_func. 보상 함수는 각 에이전트마다 하나의 완료 목록을 받습니다. actor 주문. 만약에 batch_items 키워드, trainers 또한 작업 항목도 전달해야 합니다. 공동 응답마다 하나의 보상을 반환합니다. trainers 여기서 사용하면, 해당 항목이 사용 가능한 상태에서 한 번의 작업 항목에 대해 호출합니다. batch_items[0].

reward_processor는 훈련 전에 반환된 보상을 변환합니다. 예를 들어, RewardProcessors.scale(factor=10.0)는 comlrl.utils.reward_processor에서 보상을 10배 곱합니다. 기본적으로 보상은 변경되지 않습니다.

규칙 기반 보상#

규칙 기반 보상은 응답과 작업 항목을 직접 사용하여 점수를 계산합니다. 예를 들어, 작업 검사, 검증기 또는 테스트 함을 사용합니다. 다음 함수는 길이 비율 작업의 점수를 계산합니다:

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

이 함수는 정렬된 응답 쌍당 하나의 공유 보상을 반환합니다. 작업은 문자 수 비율을 보상합니다. 이는 작성 품질을 측정하지 않으며, 에이전트는 독립적으로 생성합니다.

사용자는 모델을 로드하지 않고 작업 세트, 포맷터 및 보상 함수를 확인할 수 있습니다:

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]

LLM 기반 보상 모델#

보상 함수는 기존 LLM 보상 모델과 응답을 점수할 수도 있습니다. trainers는 반환된 값만 사용합니다. 다음 스케치는 보상 모델을 단일 점수 분류 헤드로 감싸고 있습니다. reward_model_id는 자리 표시자입니다:

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,

보상 모델은 한 번 로드하고, 함수 외부에서 장치에 배치하여 메모리가 충분한 경우, trainers를 통해 이동하거나 업데이트하지 않습니다. 선호도를 기반으로 보상 모델을 학습하려면 LLM 보상 모델을 참조하세요.

훈련 및 평가#

사용자는 train_dataset 및 eval_dataset을 통해 개별적으로 훈련 및 평가 작업 항목을 제공할 수 있습니다. 평가에는 훈련 형식 및 보상 함수가 사용되며 모델 가중치를 업데이트하지 않습니다. 출력은 구성된 생성 설정으로 샘플링되므로 반복적인 평가는 다른 점수를 생성할 수 있습니다.

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)

MAREINFORCE, IAC 및 MAAC 패밀리에 대해, eval_interval는 훈련 작업 항목을 계산합니다. 평가는 각 epoch의 0번째 항목 전에 실행되고, 그런 다음 구성된 간격으로 실행됩니다; 0는 수동 평가만 활성화합니다. eval_num_samples는 평가된 작업 항목에 대한 양의 제한을 지정하며, 평가 분의 시작 부분부터 시작합니다. 이를 분의 길이에 설정하면 모든 작업 항목을 평가합니다. eval_batch_size는 에이전트당 생성되는 응답 수를 제어하는 평가 데이터 로더 배치 크기를 제어하며, 에이전트당 생성되는 응답 수를 제어하지 않습니다.

MAGRPO, IAC 및 MAAC는 평가 중에 에이전트당 턴당 하나의 응답을 생성합니다. IAC 및 MAAC는 구성된 턴 구조와 함께 rollout 수집 경로를 사용하지만, evaluate() 메서드는 샘플 수 인수를 사용하지 않습니다:

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

eval_dataset가 없는 경우, evaluate()는 빈 딕셔너리를 반환합니다. PbRL 및 trainers에 대해, 수동 evaluate(num_eval_samples=...)도 사용할 수 있습니다; 자동 평가 일정은 위에서 언급한 작업 항목 일정 대신 선호 학습 또는 actor 훈련 단계를 따릅니다.

로깅#

CoMLRL는 Weights & Biases에서 훈련 곡선, 작업별 지표 및 로컬 평가 기록을 지원합니다. 예제에서는 trainer 및 위에서 언급한 평가 분을 사용합니다.

Weights & Biases#

사용자는 trainer를 만들 때 wandb_config를 제공하여 Weights & Biases에서 훈련 곡선을 기록할 수 있습니다:

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는 로깅이 활성화되기 전에 사용자의 환경에서 인증할 수 있습니다. wandb_config 또는 None로 설정하면 W&B가 비활성화됩니다. 지원되는 런 설정에는 project, entity, name (또는 run_name), tags, 및 dir가 포함됩니다. config_sections 필드는 기록된 구성에 작업 및 실험 설정을 추가합니다; config_sections.output.verbose는 trainer 진행 출력 제어합니다.

학습 지표#

Trainers 턴당 한 번씩 기본적으로 다음 지표를 기록 k (1)부터 시작. 사용자 정의 지표 그 옆에 함께 작업 관련 값을 추가할 수 있습니다.

지표	의미
`turn_k/reward_mean`	훈련 rollouts의 턴 `k`에 대한 평균 보상
`turn_k/expected_return`	훈련 rollouts의 턴 `k`부터의 평균 반환
`eval/turn_k/reward_mean`	평가 rollouts의 턴 `k`에 대한 평균 보상
`eval/turn_k/expected_return`	평가 rollouts의 턴 `k`부터의 평균 반환

참조 KL가 활성화된 경우, reference_kl_mean 및 reference_kl_penalty_mean도 각 턴에 기록됩니다. 훈련 지표는 env_step, 즉 공동 행동 보상 평가 횟수에 대해 기록됩니다. 이는 작업 인덱스보다 빠르게 증가합니다 (여러 세대 또는 턴의 경우). logging_steps는 훈련 로그 간의 최소 env_step 간격을 설정합니다. 평가 지표는 eval_interval를 따릅니다.

사용자 지정 지표#

사용자는 MAGRPO에 사용자 지정 평가 지표를 추가하기 위해 eval_logger와 eval_aggregator를 함께 제공할 수 있습니다. 로거는 [agent][task item][turn] 순서로 완료를 수신하고 원래 batch_items와 원시 prompts를 수신합니다. 어그리게이터는 eval/ 접두사가 없는 스칼라 지표를 반환합니다. trainer는 이를 추가합니다.

다음 예는 각 턴에서 에이전트 간의 평균 응답 길이를 측정합니다:

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,

결과적으로 생성된 딕셔너리에는 eval/turn_1/response_chars 가 포함됩니다. 이러한 콜백 함수 또한 작업-아이템 필드를 수신하여, 프롬프트 내에 참조 답변 없이 작업별 정확도 또는 그룹 메트릭을 계산할 수 있도록 합니다.

IAC 및 MAAC 대신 metrics_callback(rollouts) 를 수신하며, 이는 숫자 값의 딕셔너리를 반환합니다. 각 rollout 샘플은 다음과 같은 필드를 제공합니다: agent_idx, prompt, completion, 및 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(...).

이 콜백 함수는 학습 및 평가 시 모두 실행됩니다. 평가 시에는 평가된 작업 아이템의 값을 평균하여 eval/ 를 추가합니다. 학습 시에는 rollout 수집 시간에 별도로 표준 logging_steps 게이트에서 값을 기록합니다.

지역 메트릭#

평가는 W&B 가 비활성화되어 있더라도 메트릭 딕셔너리를 반환합니다. 이 딕셔너리는 완료된 평가를 기록하기 위해 로컬에 저장할 수 있습니다:

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