从第一原理构建 AI Agent

class Canvas:
    def __init__(self):
        self.shapes = {}      # shape_id -> Shape 对象
        self.connections = []  # [(from_id, to_id), ...]

    def step(self, action: dict) -> dict:
        """接收模型输出的 JSON action，返回更新后的状态"""
        action_type = action.get("type")
        if action_type == "create_shape":
            shape_id = action["shape_id"]
            shape_type = action["shape_type"]  # "circle" | "square" | "label"
            x, y = action["x"], action["y"]
            label = action.get("label", "")
            self.shapes[shape_id] = Shape(shape_id, shape_type, x, y, label)
            return {"status": "ok", "shapes": len(self.shapes)}
        elif action_type == "connect":
            self.connections.append((action["from_id"], action["to_id"]))
            return {"status": "ok", "connections": len(self.connections)}
        else:
            return {"status": "error", "msg": f"unknown action: {action_type}"}

def compute_reward(canvas: Canvas, user_prompt: str) -> dict:
    r = {}

    # 1. 语法正确性：所有 action 都成功了吗？
    r["syntax"] = 1.0 if canvas.errors == 0 else -1.0

    # 2. 布局质量：形状之间不重叠
    overlaps = 0
    for a in canvas.shapes:
        for b in canvas.shapes:
            if a != b and bounding_box_overlap(a, b):
                overlaps += 1
    r["layout"] = max(0, 1.0 - 0.2 * overlaps)

    # 3. 语义覆盖：label 包含 prompt 中的关键词
    keywords = extract_keywords(user_prompt)
    matched = sum(
        1 for kw in keywords
        if kw in [s.label.lower() for s in canvas.shapes.values()]
    )
    r["semantic"] = matched / max(len(keywords), 1)

    r["total"] = 0.3 * r["syntax"] + 0.3 * r["layout"] + 0.4 * r["semantic"]
    return r

def generate_teacher_trajectories(prompts: list[str]) -> list[dict]:
    trajectories = []
    for prompt in prompts:
        # 1. 用强模型生成 action 序列
        teacher_output = gemini.generate(
            f"User request: {prompt}\n"
            "Output a sequence of JSON actions to create the diagram."
        )
        # 2. 在环境中验证
        canvas = Canvas()
        for action in parse_actions(teacher_output):
            result = canvas.step(action)
            if result["status"] == "error":
                break  # 这条轨迹无效
        else:
            # 3. 全部执行成功 → 保留
            reward = compute_reward(canvas, prompt)
            if reward["total"] > 0.7:  # 质量阈值
                trajectories.append({
                    "prompt": prompt,
                    "actions": teacher_output,
                    "reward": reward["total"]
                })
    return trajectories

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")

# 只对 action 部分计算 loss，忽略 prompt
collator = DataCollatorForCompletionOnlyLM(
    response_template="\n[ACTIONS]",
    tokenizer=tokenizer
)

trainer = SFTTrainer(
    model=model,
    train_dataset=teacher_trajectories,
    args=TrainingArguments(
        output_dir="./sft_agent",
        per_device_train_batch_size=4,
        num_train_epochs=3,
        learning_rate=2e-5,
    ),
    data_collator=collator,
)
trainer.train()

def compute_advantages(rewards: list[float], eps=1e-8):
    """Group-relative advantage: A_i = (r_i - mean(r)) / (std(r) + eps)"""
    mean_r = sum(rewards) / len(rewards)
    std_r = (sum((r - mean_r) ** 2 for r in rewards) / len(rewards)) ** 0.5
    return [(r - mean_r) / (std_r + eps) for r in rewards]

for step in range(num_steps):
    batch = []
    # 1. 采样：对每个 prompt 生成 N 个 action 序列
    for prompt in prompts:
        for _ in range(group_size):
            actions = model.generate(prompt)
            canvas = Canvas()
            for a in parse_actions(actions):
                canvas.step(a)
            reward = compute_reward(canvas, prompt)
            batch.append((prompt, actions, reward))

    # 2. 计算 group-relative advantage
    rewards = [item[2] for item in batch]
    advantages = compute_advantages(rewards)

    # 3. 策略梯度更新
    for (prompt, actions, _), advantage in zip(batch, advantages):
        log_prob = model.log_prob(prompt, actions)
        loss = -log_prob * advantage  # 高 advantage → 提高概率
        loss.backward()

    optimizer.step()
    optimizer.zero_grad()