この例では、2つのエージェントがタスクに回答し、回答に対してフィードバックを受け、その後再度回答します。
設定#
このタスクでは、2番目のエージェントに、最初のエージェントの回答よりも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 は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つの共同回答を形成します。
- フィードバック関数は、各共同回答を新しいプロンプトに変換します。
- 2回目のターンで、各エージェントが再度回答するため、1つのタスクには最大16回の共同応答が存在します。
- 1回目のターンでの回答は、自身の報酬に加えて、それに続く2回目のターンでの回答の平均報酬で評価されます。
- W&Bを有効にした場合、
turn_1/reward_meanとturn_2/reward_meanが、2回目の回答が優れているかどうかを示します。