在此示例中，两个智能体回答任务，接收对答案的反馈，然后再次回答。
设置#
任务要求第二个智能体写出的答案长度为第一个智能体的答案长度的 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)
]两轮训练#
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显示第二回合的回答是否更好。