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

chapter6/rgb-sim2real-grasping/pipeline.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""Experiment 6-13: RGB domain-transfer benchmark on the local GPU.
+
+The benchmark is intentionally self-contained.  It trains a small RGB policy
+on a source visual domain and evaluates it on a shifted target domain, with
+and without domain randomization.  It is a production-grade local proxy for
+the sim-to-real argument; SO-100 deployment remains a separately gated
+extension and is never implied by this run.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+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
+
+IMAGE_SIZE = 32
+CLASSES = ("left", "right", "up", "down", "grasp")
+
+
+def label_for(object_xy: tuple[float, float], target_xy: tuple[float, float]) -> int:
+    dx, dy = target_xy[0] - object_xy[0], target_xy[1] - object_xy[1]
+    if abs(dx) < 0.10 and abs(dy) < 0.10:
+        return 4
+    if abs(dx) >= abs(dy):
+        return 1 if dx > 0 else 0
+    return 3 if dy > 0 else 2
+
+
+def render_sample(rng: np.random.Generator, domain: str) -> tuple[np.ndarray, int]:
+    object_xy = tuple(rng.uniform(0.18, 0.82, size=2))
+    target_xy = tuple(rng.uniform(0.18, 0.82, size=2))
+    label = label_for(object_xy, target_xy)
+    if domain == "source_clean":
+        background = np.full((IMAGE_SIZE, IMAGE_SIZE, 3), [220, 220, 220], dtype=np.float32)
+        object_color, target_color = (214, 60, 60), (45, 130, 65)
+        noise = 0.0
+    elif domain in {"source_background", "source_full"}:
+        if domain == "source_background":
+            base = rng.uniform(45, 225, size=(1, 1, 3)).astype(np.float32)
+            background = np.broadcast_to(base, (IMAGE_SIZE, IMAGE_SIZE, 3)).copy()
+            background += rng.normal(0, 9, size=background.shape)
+            object_color, target_color = (214, 60, 60), (45, 130, 65)
+        else:
+            base = rng.uniform(45, 225, size=(1, 1, 3)).astype(np.float32)
+            background = np.broadcast_to(base, (IMAGE_SIZE, IMAGE_SIZE, 3)).copy()
+            background += rng.normal(0, 9, size=background.shape)
+            object_color = (int(rng.uniform(180, 245)), int(rng.uniform(45, 105)), int(rng.uniform(35, 100)))
+            target_color = (int(rng.uniform(35, 100)), int(rng.uniform(125, 205)), int(rng.uniform(75, 170)))
+        noise = 3.0
+    elif domain == "source_appearance":
+        background = np.full((IMAGE_SIZE, IMAGE_SIZE, 3), rng.uniform(150, 235), dtype=np.float32)
+        object_color = (int(rng.uniform(170, 245)), int(rng.uniform(45, 110)), int(rng.uniform(35, 100)))
+        target_color = (int(rng.uniform(35, 100)), int(rng.uniform(120, 210)), int(rng.uniform(75, 175)))
+        noise = 2.0
+    elif domain in {"target_realistic", "target_bright"}:
+        base = rng.uniform(45, 175, size=(IMAGE_SIZE, IMAGE_SIZE, 1))
+        stripes = (np.sin(np.arange(IMAGE_SIZE)[None, :, None] / 3.0) * 18.0).astype(np.float32)
+        background = np.repeat(base, 3, axis=2) + stripes
+        if domain == "target_bright":
+            background = np.clip(background + 55, 0, 255)
+        object_color, target_color = (235, 185, 55), (70, 160, 205)
+        noise = 14.0
+    else:
+        raise ValueError(f"unknown domain: {domain}")
+    image = Image.fromarray(np.uint8(np.clip(background, 0, 255)), mode="RGB")
+    draw = ImageDraw.Draw(image)
+    ox, oy = int(object_xy[0] * IMAGE_SIZE), int(object_xy[1] * IMAGE_SIZE)
+    tx, ty = int(target_xy[0] * IMAGE_SIZE), int(target_xy[1] * IMAGE_SIZE)
+    draw.rectangle((tx - 4, ty - 4, tx + 4, ty + 4), outline=target_color, width=2)
+    draw.ellipse((ox - 3, oy - 3, ox + 3, oy + 3), fill=object_color, outline=(20, 20, 20))
+    array = np.asarray(image, dtype=np.float32)
+    if noise:
+        array += rng.normal(0, noise, size=array.shape)
+    array = np.clip(array / 255.0, 0.0, 1.0).transpose(2, 0, 1)
+    return array.astype(np.float32), label
+
+
+def make_dataset(size: int, domain: str, seed: int) -> tuple[torch.Tensor, torch.Tensor]:
+    rng = np.random.default_rng(seed)
+    images, labels = zip(*(render_sample(rng, domain) for _ in range(size)))
+    return torch.from_numpy(np.stack(images)), torch.tensor(labels, dtype=torch.long)
+
+
+class RGBPolicy(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 48, 3, padding=1), nn.ReLU(),
+        )
+        self.head = nn.Sequential(nn.Flatten(), nn.Linear(48 * 8 * 8, 128), nn.ReLU(), nn.Linear(128, len(CLASSES)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.head(self.net(x))
+
+
+def train_policy(images: torch.Tensor, labels: torch.Tensor, device: torch.device, seed: int, epochs: int) -> tuple[RGBPolicy, list[float]]:
+    generator = torch.Generator().manual_seed(seed)
+    loader = DataLoader(TensorDataset(images, labels), batch_size=128, shuffle=True, generator=generator)
+    model = RGBPolicy().to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)
+    criterion = nn.CrossEntropyLoss()
+    losses: list[float] = []
+    for _ in range(epochs):
+        total, count = 0.0, 0
+        for batch_images, batch_labels in loader:
+            optimizer.zero_grad(set_to_none=True)
+            loss = criterion(model(batch_images.to(device)), batch_labels.to(device))
+            loss.backward()
+            optimizer.step()
+            total += float(loss.item()) * len(batch_labels)
+            count += len(batch_labels)
+        losses.append(total / count)
+    return model, losses
+
+
+def accuracy(model: RGBPolicy, images: torch.Tensor, labels: torch.Tensor, device: torch.device) -> float:
+    with torch.no_grad():
+        prediction = model(images.to(device)).argmax(dim=1).cpu()
+    return float((prediction == labels).float().mean().item())
+
+
+def save_preview(path: Path, images: torch.Tensor, labels: torch.Tensor) -> None:
+    tile = (images[:16].permute(0, 2, 3, 1).numpy() * 255).astype(np.uint8)
+    canvas = Image.new("RGB", (IMAGE_SIZE * 4, IMAGE_SIZE * 4), "white")
+    for index, array in enumerate(tile):
+        canvas.paste(Image.fromarray(array), ((index % 4) * IMAGE_SIZE, (index // 4) * IMAGE_SIZE))
+    path.parent.mkdir(parents=True, exist_ok=True)
+    canvas.save(path)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train-size", type=int, default=4096, help="examples per seed and training variant")
+    parser.add_argument("--test-size", type=int, default=1024, help="examples per seed and target domain")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--seeds", default="20260808,20260809,20260810")
+    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(",")]
+    except ValueError:
+        parser.error("--seeds must be comma-separated integers")
+    if args.train_size < 2048 or args.test_size < 512 or args.epochs < 8 or len(seeds) < 3:
+        parser.error("the benchmark requires at least 2048 train examples, 512 test examples, 8 epochs and 3 seeds")
+    seed_everything(seeds[0])
+    try:
+        device = select_device(not args.allow_cpu)
+    except RuntimeError as exc:
+        parser.error(str(exc))
+    started = time.perf_counter()
+    variants = ("source_clean", "source_background", "source_appearance", "source_full")
+    target_domains = ("target_realistic", "target_bright")
+    rows: list[dict[str, Any]] = []
+    checkpoint_model: RGBPolicy | None = None
+    for seed in seeds:
+        source_test_x, source_test_y = make_dataset(args.test_size, "source_clean", seed + 100)
+        target_tests = {domain: make_dataset(args.test_size, domain, seed + 200 + index) for index, domain in enumerate(target_domains)}
+        for variant_index, variant in enumerate(variants):
+            train_x, train_y = make_dataset(args.train_size, variant, seed + variant_index)
+            model, losses = train_policy(train_x, train_y, device, seed + 50 + variant_index, args.epochs)
+            if seed == seeds[0] and variant == "source_full":
+                checkpoint_model = model
+            rows.append({"seed": seed, "variant": variant, "source_accuracy": accuracy(model, source_test_x, source_test_y, device), "target_accuracy": {domain: accuracy(model, images, labels, device) for domain, (images, labels) in target_tests.items()}, "final_loss": losses[-1]})
+    grouped: dict[str, dict[str, Any]] = {}
+    for variant in variants:
+        grouped[variant] = {}
+        for domain in target_domains:
+            values = [row["target_accuracy"][domain] for row in rows if row["variant"] == variant]
+            grouped[variant][domain] = {"mean": float(np.mean(values)), "std": float(np.std(values)), "values": values}
+        source_values = [row["source_accuracy"] for row in rows if row["variant"] == variant]
+        grouped[variant]["source"] = {"mean": float(np.mean(source_values)), "std": float(np.std(source_values)), "values": source_values}
+    replay_a = make_dataset(256, "source_full", seeds[0])
+    replay_b = make_dataset(256, "source_full", seeds[0])
+    metrics = {"device": device_info(device), "protocol": {"seeds": seeds, "variants": list(variants), "target_domains": list(target_domains), "train_size_per_variant": args.train_size, "test_size_per_domain": args.test_size, "epochs": args.epochs, "total_training_examples": len(seeds) * len(variants) * args.train_size}, "rows": rows, "summary": grouped, "dataset_replay_match": bool(torch.equal(replay_a[0], replay_b[0]) and torch.equal(replay_a[1], replay_b[1])), "wall_time_ms": round((time.perf_counter() - started) * 1000, 3)}
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    preview_source = args.output_dir / "source_preview.png"
+    preview_target = args.output_dir / "target_preview.png"
+    source_preview_x, source_preview_y = make_dataset(64, "source_clean", seeds[0] + 1000)
+    target_preview_x, target_preview_y = make_dataset(64, "target_realistic", seeds[0] + 1001)
+    save_preview(preview_source, source_preview_x, source_preview_y)
+    save_preview(preview_target, target_preview_x, target_preview_y)
+    checkpoint_path = args.output_dir / "randomized_policy.pt"
+    if checkpoint_model is None:
+        raise RuntimeError("source_full checkpoint was not produced")
+    torch.save({"model": {key: value.detach().cpu() for key, value in checkpoint_model.state_dict().items()}, "seed": seeds[0], "classes": CLASSES}, checkpoint_path)
+    matrix_path = args.output_dir / "matrix.json"
+    write_json(matrix_path, {"rows": rows, "summary": grouped})
+    metrics_path = args.output_dir / "metrics.json"
+    write_json(metrics_path, metrics)
+    evidence = {"schema_version": "3.0", "experiment_id": "6-13", "status": "complete", "kind": "local_gpu_rgb_domain_transfer", "seed": seeds[0], "metrics": metrics, "artifacts": [{"kind": "checkpoint", "path": relative_or_absolute(checkpoint_path, args.output_dir), "sha256": sha256(checkpoint_path)}, {"kind": "metrics", "path": relative_or_absolute(metrics_path, args.output_dir), "sha256": sha256(metrics_path)}, {"kind": "matrix", "path": relative_or_absolute(matrix_path, args.output_dir), "sha256": sha256(matrix_path)}, {"kind": "source_preview", "path": relative_or_absolute(preview_source, args.output_dir), "sha256": sha256(preview_source)}, {"kind": "target_preview", "path": relative_or_absolute(preview_target, args.output_dir), "sha256": sha256(preview_target)}], "hardware_extension": {"status": "gated", "actuation_attempted": False, "upstream": "StoneT2000/lerobot-sim2real"}, "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())