Initial pi0.5 minimal validation

2026-07-11 20:04:18 +08:00
commit ffbcca8d2b
20 changed files with 2354 additions and 0 deletions
@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .data import Pi05OverfitDataset
+from .flow import FlowMatching
+from .model import Pi05Config, Pi05Mini
+from .runtime import choose_device, write_json
+
+
+@torch.no_grad()
+def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    cfg = checkpoint["config"]
+    root = Path(__file__).resolve().parents[2]
+    device = choose_device(cfg.get("device", "auto"))
+    data_manifest = Path(manifest) if manifest else root / "data/overfit/overfit_manifest.jsonl"
+    dataset = Pi05OverfitDataset(
+        data_manifest, image_size=cfg["data"]["image_size"],
+        action_horizon=cfg["data"]["action_horizon"],
+        action_statistics=checkpoint["action_statistics"]
+    )
+    model = Pi05Mini(Pi05Config(**checkpoint["model_config"]))
+    model.load_state_dict(checkpoint["model"])
+    model.to(device).eval()
+    flow = FlowMatching(cfg["inference"]["flow_steps"])
+    rows = []
+    figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False)
+    for index, sample in enumerate(dataset):
+        image = sample["image"][None].to(device)
+        tokens = sample["prefix_tokens"][None].to(device)
+        # Autoregressively predict semantic subtask followed by the complete
+        # discretized action stream (ACTION_START, 32 scalar tokens, ACTION_END).
+        zero_action = torch.zeros(1, dataset.action_horizon, 2, device=device)
+        generated: list[int] = []
+        expected_tokens = 1 + 2 + dataset.action_horizon * 2
+        current_tokens = tokens
+        for _ in range(expected_tokens):
+            semantic = model(
+                image, current_tokens, zero_action,
+                torch.zeros(1, device=device),
+                generated_token_start=tokens.shape[1]
+            )
+            next_token = int(semantic.text_logits[0, -1].argmax().cpu())
+            generated.append(next_token)
+            current_tokens = torch.cat(
+                (current_tokens, torch.tensor([[next_token]], device=device)), dim=1
+            )
+        predicted_token = generated[0]
+        vlm_action_tokens = torch.tensor(generated[1:], dtype=torch.long)
+        try:
+            vlm_normalized = dataset.tokenizer.decode_actions(vlm_action_tokens)
+            vlm_xy = dataset.action_normalizer.decode(
+                vlm_normalized, sample["history_xy"]
+            ).numpy()
+            vlm_valid = True
+        except ValueError:
+            vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
+            vlm_valid = False
+        conditioned_tokens = torch.cat(
+            (tokens, torch.tensor([[predicted_token]], device=device)), dim=1
+        )
+        def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+            return model(
+                image, conditioned_tokens, x, t,
+                generated_token_start=tokens.shape[1]
+            ).vector_field
+        normalized = flow.sample(
+            velocity, (1, dataset.action_horizon, 2), device,
+            steps=cfg["inference"]["flow_steps"],
+            generator=None,
+        )[0].cpu()
+        predicted_xy = dataset.action_normalizer.decode(
+            normalized, sample["history_xy"]
+        ).numpy()
+        target_xy = sample["future_xy"].numpy()
+        diffusion_mse = float(np.mean((predicted_xy - target_xy) ** 2))
+        vlm_mse = float(np.mean((vlm_xy - target_xy) ** 2)) if vlm_valid else None
+        target_label = dataset.tokenizer.subtasks.decode(sample["subtask_target"])
+        try:
+            predicted_label = dataset.tokenizer.subtasks.decode(predicted_token)
+        except ValueError:
+            predicted_label = f"token_{predicted_token}"
+        rows.append({
+            "sample_id": sample["sample_id"], "target_subtask": target_label,
+            "predicted_subtask": predicted_label,
+            "vlm_action_tokens_valid": vlm_valid,
+            "vlm_trajectory_mse": vlm_mse,
+            "diffusion_trajectory_mse": diffusion_mse,
+        })
+        ax = axes[0, index]
+        ax.plot(target_xy[:, 0], target_xy[:, 1], "o-", label="target")
+        ax.plot(vlm_xy[:, 0], vlm_xy[:, 1], "s--", label="VLM tokens")
+        ax.plot(predicted_xy[:, 0], predicted_xy[:, 1], "x--", label="diffusion")
+        ax.scatter([0], [0], c="black", marker="s", label="ego")
+        vlm_text = "invalid" if vlm_mse is None else f"{vlm_mse:.3f}"
+        ax.set_title(
+            f"{predicted_label}\nVLM={vlm_text}, diffusion={diffusion_mse:.3f}"
+        )
+        ax.set_xlabel("forward x (m)")
+        ax.set_ylabel("left y (m)")
+        ax.axis("equal")
+        ax.grid(True, alpha=.3)
+        ax.legend(fontsize=8)
+    figure.tight_layout()
+    output_dir = Path(checkpoint_path).parent
+    plot_path = output_dir / "trajectory_comparison.png"
+    figure.savefig(plot_path, dpi=160)
+    plt.close(figure)
+    write_json(output_dir / "evaluation.json", rows)
+    print(rows)
+    return plot_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt")
+    parser.add_argument("--manifest")
+    args = parser.parse_args()
+    evaluate(args.checkpoint, args.manifest)
+
+
+if __name__ == "__main__":
+    main()