ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter6/gemini-xlerobot-navigation/desktop_planner.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+"""Experiment 6-12: desktop manipulation planning with a local GPU backend.
+
+The local run keeps the RoboCrew-style tool contract and the XLeRobot adapter
+boundary, but executes against a deterministic tabletop simulator.  This
+makes planner, postcondition, retry and short-horizon world-model behavior
+fully reproducible without claiming that a real robot moved.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+from torch import nn
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from robotics_lab_common import device_info, relative_or_absolute, select_device, seed_everything, sha256, write_json
+
+OBJECTS = ("red_cup", "yellow_paper")
+TARGETS = ("tray", "bin")
+ACTIONS = ("pick_red_cup", "place_red_cup", "pick_yellow_paper", "place_yellow_paper")
+TOOL_NAMES = ("observe_scene", "pick", "place", "verify_state", "stop")
+
+
+@dataclass
+class DesktopState:
+    status: list[int]  # 0=on table, 1=held, 2=placed
+    target_available: list[bool]
+
+    def copy(self) -> "DesktopState":
+        return DesktopState(self.status[:], self.target_available[:])
+
+    def done(self) -> bool:
+        return self.status == [2, 2]
+
+
+class DesktopToolAdapter:
+    """RoboCrew-compatible semantic tool boundary for local validation."""
+
+    def __init__(self, seed: int, failure_probability: float = 0.25):
+        self.rng = random.Random(seed)
+        self.state = DesktopState([0, 0], [True, True])
+        self.failure_probability = failure_probability
+        self.injected_failure = False
+        self.events: list[dict[str, Any]] = []
+
+    def observe_scene(self) -> dict[str, Any]:
+        observation = {"objects": dict(zip(OBJECTS, self.state.status)), "targets_available": dict(zip(TARGETS, self.state.target_available))}
+        self.events.append({"tool": "observe_scene", "ok": True, "observation": observation})
+        return observation
+
+    def _maybe_fail(self, action: str) -> bool:
+        if action == "pick_yellow_paper" and not self.injected_failure and self.rng.random() < self.failure_probability:
+            self.injected_failure = True
+            return True
+        return False
+
+    def pick(self, object_id: str) -> dict[str, Any]:
+        if object_id not in OBJECTS:
+            result = {"ok": False, "reason": "unknown_object"}
+        else:
+            index = OBJECTS.index(object_id)
+            action = f"pick_{object_id}"
+            if self._maybe_fail(action):
+                result = {"ok": False, "reason": "injected_transient_grasp_failure"}
+            elif self.state.status[index] != 0:
+                result = {"ok": False, "reason": "object_not_on_table"}
+            else:
+                self.state.status[index] = 1
+                result = {"ok": True, "postcondition": f"{object_id}=held"}
+        self.events.append({"tool": "pick", "object_id": object_id, **result})
+        return result
+
+    def place(self, object_id: str, target_id: str) -> dict[str, Any]:
+        if object_id not in OBJECTS or target_id not in TARGETS:
+            result = {"ok": False, "reason": "unknown_object_or_target"}
+        else:
+            oi, ti = OBJECTS.index(object_id), TARGETS.index(target_id)
+            if self.state.status[oi] != 1:
+                result = {"ok": False, "reason": "object_not_held"}
+            elif not self.state.target_available[ti]:
+                result = {"ok": False, "reason": "target_unavailable"}
+            else:
+                self.state.status[oi] = 2
+                self.state.target_available[ti] = False
+                result = {"ok": True, "postcondition": f"{object_id}=in_{target_id}"}
+        self.events.append({"tool": "place", "object_id": object_id, "target_id": target_id, **result})
+        return result
+
+    def verify_state(self) -> dict[str, Any]:
+        result = {"ok": self.state.done(), "state": self.observe_scene()}
+        self.events.append({"tool": "verify_state", **result})
+        return result
+
+    def stop(self) -> dict[str, Any]:
+        result = {"ok": True, "stopped": True}
+        self.events.append({"tool": "stop", **result})
+        return result
+
+
+def encode_state(state: DesktopState) -> list[float]:
+    return [state.status[0] / 2.0, state.status[1] / 2.0, float(state.target_available[0]), float(state.target_available[1])]
+
+
+def transition(state: DesktopState, action: int) -> DesktopState:
+    next_state = state.copy()
+    if action == 0 and next_state.status[0] == 0:
+        next_state.status[0] = 1
+    elif action == 1 and next_state.status[0] == 1 and next_state.target_available[0]:
+        next_state.status[0], next_state.target_available[0] = 2, False
+    elif action == 2 and next_state.status[1] == 0:
+        next_state.status[1] = 1
+    elif action == 3 and next_state.status[1] == 1 and next_state.target_available[1]:
+        next_state.status[1], next_state.target_available[1] = 2, False
+    return next_state
+
+
+class WorldModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 4))
+
+    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
+        return self.net(torch.cat([state, action], dim=-1))
+
+
+def train_world_model(device: torch.device, seed: int, epochs: int = 160) -> tuple[WorldModel, float, float]:
+    rng = random.Random(seed)
+    states, actions, targets = [], [], []
+    for _ in range(12000):
+        state = DesktopState([rng.randrange(3), rng.randrange(3)], [bool(rng.randrange(2)), bool(rng.randrange(2))])
+        action_index = rng.randrange(len(ACTIONS))
+        nxt = transition(state, action_index)
+        states.append(encode_state(state))
+        one_hot = [1.0 if index == action_index else 0.0 for index in range(len(ACTIONS))]
+        actions.append(one_hot)
+        targets.append(encode_state(nxt))
+    x_state = torch.tensor(states, dtype=torch.float32, device=device)
+    x_action = torch.tensor(actions, dtype=torch.float32, device=device)
+    y = torch.tensor(targets, dtype=torch.float32, device=device)
+    split = int(len(y) * 0.8)
+    model = WorldModel().to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)
+    loss_fn = nn.MSELoss()
+    for _ in range(epochs):
+        optimizer.zero_grad(set_to_none=True)
+        loss = loss_fn(model(x_state[:split], x_action[:split]), y[:split])
+        loss.backward()
+        optimizer.step()
+    with torch.no_grad():
+        train_loss = float(loss_fn(model(x_state[:split], x_action[:split]), y[:split]).item())
+        test_loss = float(loss_fn(model(x_state[split:], x_action[split:]), y[split:]).item())
+    return model, train_loss, test_loss
+
+
+def action_from_index(index: int) -> tuple[str, str | None, str | None]:
+    mapping = [("pick", "red_cup", None), ("place", "red_cup", "tray"), ("pick", "yellow_paper", None), ("place", "yellow_paper", "bin")]
+    return mapping[index]
+
+
+def execute(adapter: DesktopToolAdapter, index: int) -> dict[str, Any]:
+    kind, object_id, target_id = action_from_index(index)
+    if kind == "pick":
+        return adapter.pick(object_id or "")
+    return adapter.place(object_id or "", target_id or "")
+
+
+def render_scene(path: Path, state: DesktopState) -> None:
+    image = Image.new("RGB", (480, 300), (235, 232, 220))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle((40, 40, 440, 260), outline=(60, 60, 60), width=3)
+    locations = [(150, 130), (250, 130)]
+    colors = [(210, 60, 60), (220, 190, 40)]
+    for idx, (x, y) in enumerate(locations):
+        if state.status[idx] != 2:
+            draw.ellipse((x - 22, y - 22, x + 22, y + 22), fill=colors[idx], outline=(20, 20, 20))
+        draw.text((x - 35, y + 30), OBJECTS[idx], fill=(20, 20, 20))
+    draw.rectangle((320, 80, 390, 145), outline=(40, 100, 210), width=3)
+    draw.rectangle((320, 170, 390, 235), outline=(40, 130, 60), width=3)
+    draw.text((325, 95), "tray", fill=(20, 20, 20))
+    draw.text((325, 185), "bin", fill=(20, 20, 20))
+    path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(path)
+
+
+def run_episode(mode: str, seed: int, model: WorldModel | None, device: torch.device, failure_probability: float) -> dict[str, Any]:
+    adapter = DesktopToolAdapter(seed, failure_probability=failure_probability)
+    if mode == "open_loop":
+        plan = [0, 1, 2, 3]
+        for action in plan:
+            execute(adapter, action)
+    else:
+        for _ in range(12):
+            if adapter.state.done():
+                break
+            if mode == "predictive" and model is not None:
+                state_tensor = torch.tensor([encode_state(adapter.state)], dtype=torch.float32, device=device)
+                candidates = [index for index in range(4) if (index in (0, 2) and adapter.state.status[index // 2] == 0) or (index in (1, 3) and adapter.state.status[index // 2] == 1)]
+                if not candidates:
+                    break
+                action_vectors = torch.eye(4, device=device)[candidates]
+                with torch.no_grad():
+                    predicted = model(state_tensor.repeat(len(candidates), 1), action_vectors)
+                score = predicted[:, 0] + predicted[:, 1] + (predicted[:, 0] > 0.95).float() + (predicted[:, 1] > 0.95).float()
+                action = candidates[int(torch.argmax(score).item())]
+            else:
+                action = next((index for index in (0, 1, 2, 3) if (index in (0, 2) and adapter.state.status[index // 2] == 0) or (index in (1, 3) and adapter.state.status[index // 2] == 1)), 0)
+            result = execute(adapter, action)
+            if not result.get("ok"):
+                adapter.observe_scene()
+        adapter.verify_state()
+    return {"success": adapter.state.done(), "tool_calls": len(adapter.events), "recoveries": sum(1 for event in adapter.events if event.get("reason") == "injected_transient_grasp_failure"), "events": adapter.events}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--episodes", type=int, default=128, help="episodes per seed/failure/mode cell")
+    parser.add_argument("--seeds", default="20260808,20260809,20260810")
+    parser.add_argument("--failure-probabilities", default="0.0,0.25,0.5")
+    parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "validation" / "runs" / "local-gpu")
+    parser.add_argument("--allow-cpu", action="store_true")
+    args = parser.parse_args()
+    try:
+        seeds = [int(value) for value in args.seeds.split(",")]
+        failure_probabilities = [float(value) for value in args.failure_probabilities.split(",")]
+    except ValueError:
+        parser.error("--seeds and --failure-probabilities must be comma-separated values")
+    if args.episodes < 64 or len(seeds) < 3 or not failure_probabilities or any(value < 0 or value > 1 for value in failure_probabilities):
+        parser.error("use at least three seeds, 64 episodes per cell, and probabilities in [0,1]")
+    seed_everything(seeds[0])
+    try:
+        device = select_device(not args.allow_cpu)
+    except RuntimeError as exc:
+        parser.error(str(exc))
+    started = time.perf_counter()
+    results: list[dict[str, Any]] = []
+    model_reports: list[dict[str, Any]] = []
+    event_log_rows: list[dict[str, Any]] = []
+    for model_seed in seeds:
+        model, train_loss, test_loss = train_world_model(device, model_seed)
+        model_reports.append({"seed": model_seed, "train_mse": train_loss, "test_mse": test_loss})
+        for failure_probability in failure_probabilities:
+            for mode in ("open_loop", "closed_loop", "predictive"):
+                episodes = [run_episode(mode, model_seed + index, model if mode == "predictive" else None, device, failure_probability) for index in range(args.episodes)]
+                results.append({"seed": model_seed, "failure_probability": failure_probability, "mode": mode, "episodes": args.episodes, "success_rate": sum(item["success"] for item in episodes) / args.episodes, "mean_tool_calls": sum(item["tool_calls"] for item in episodes) / args.episodes, "recoveries": sum(item["recoveries"] for item in episodes)})
+                for episode_index, episode in enumerate(episodes):
+                    event_log_rows.append({
+                        "seed": model_seed,
+                        "failure_probability": failure_probability,
+                        "mode": mode,
+                        "episode": episode_index,
+                        "success": episode["success"],
+                        "events": episode["events"],
+                    })
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    scene_path = args.output_dir / "scene_initial.png"
+    render_scene(scene_path, DesktopState([0, 0], [True, True]))
+    events_path = args.output_dir / "predictive_episode_events.json"
+    write_json(events_path, {"tools": TOOL_NAMES, "episodes": event_log_rows})
+    replay_model, _, _ = train_world_model(device, seeds[0])
+    replay_a = run_episode("predictive", seeds[0], replay_model, device, failure_probabilities[-1])
+    replay_b = run_episode("predictive", seeds[0], replay_model, device, failure_probabilities[-1])
+    metrics = {"device": device_info(device), "protocol": {"seeds": seeds, "failure_probabilities": failure_probabilities, "episodes_per_cell": args.episodes, "total_episodes": len(results) * args.episodes}, "models": model_reports, "cells": results, "deterministic_replay": replay_a == replay_b, "wall_time_ms": round((time.perf_counter() - started) * 1000, 3)}
+    metrics_path = args.output_dir / "metrics.json"
+    write_json(metrics_path, metrics)
+    evidence = {"schema_version": "3.0", "experiment_id": "6-12", "status": "complete", "kind": "desktop_manipulation_planning", "seed": seeds[0], "tool_contract": list(TOOL_NAMES), "metrics": metrics, "artifacts": [{"kind": "metrics", "path": relative_or_absolute(metrics_path, args.output_dir), "sha256": sha256(metrics_path)}, {"kind": "events", "path": relative_or_absolute(events_path, args.output_dir), "sha256": sha256(events_path)}, {"kind": "scene", "path": relative_or_absolute(scene_path, args.output_dir), "sha256": sha256(scene_path)}], "xlerobot_robocrew_extension": {"status": "gated", "tool_adapter_required": True, "actuation_attempted": False}, "blockers": [] if not args.allow_cpu else ["CPU debug mode is not a GPU acceptance run"]}
+    evidence_path = args.output_dir / "evidence.json"
+    write_json(evidence_path, evidence)
+    print(json.dumps(evidence, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())