Initial pi0.7 minimal validation

2026-07-12 00:52:00 +08:00
commit b4c8db846b
27 changed files with 3724 additions and 0 deletions
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi07mini.data import (
+    Pi07SteerableDataset,
+    build_policy_trajectories,
+    load_manifest,
+    materialize_subset,
+)
+
+
+def _make_source_project(root: Path) -> Path:
+    processed = root / "data" / "processed"
+    (processed / "images").mkdir(parents=True)
+    (processed / "trajectories").mkdir(parents=True)
+    pixels = np.zeros((40, 64, 3), dtype=np.uint8)
+    pixels[..., 0] = 80
+    pixels[..., 1] = np.arange(64, dtype=np.uint8)[None, :]
+    Image.fromarray(pixels).save(processed / "images" / "sample.jpg")
+    hx = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
+    history = np.stack((hx, 0.01 * hx, np.zeros(8)), axis=-1)
+    fx = np.linspace(1.0, 24.0, 16, dtype=np.float32)
+    future = np.stack((fx, 0.015 * fx**2, np.zeros(16)), axis=-1)
+    np.savez_compressed(
+        processed / "trajectories" / "sample.npz",
+        history_xyz=history,
+        future_xyz=future,
+    )
+    manifest = processed / "overfit_manifest.jsonl"
+    manifest.write_text(
+        json.dumps(
+            {
+                "clip_id": "clip-a",
+                "t0_us": 123,
+                "reasoning": "Steer left.",
+                "image_path": "data/processed/images/sample.jpg",
+                "trajectory_path": "data/processed/trajectories/sample.npz",
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    return manifest
+
+
+def test_policy_variants_and_metadata() -> None:
+    x = torch.linspace(-7.0, 0.0, 8)
+    history = torch.stack((x, torch.zeros_like(x)), -1)
+    fx = torch.linspace(1.0, 24.0, 16)
+    future = torch.stack((fx, 0.02 * fx.square()), -1)
+    policies = build_policy_trajectories(history, future)
+    assert [policy.kind for policy in policies] == [
+        "expert",
+        "fast",
+        "slow",
+        "understeer",
+        "stop",
+    ]
+    metadata = {policy.kind: policy.metadata.as_dict() for policy in policies}
+    assert metadata["expert"] == {
+        "speed": "normal",
+        "quality": "expert",
+        "mistake": "none",
+        "control": "human",
+    }
+    assert metadata["understeer"]["mistake"] == "understeer"
+    assert metadata["stop"]["quality"] == "poor"
+    assert policies[0].metadata.as_token_metadata() == {
+        "speed_steps": 1_000,
+        "quality": 5,
+        "mistake": False,
+        "control_mode": "ee",
+        "source_type": "demonstration",
+    }
+    assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4])
+
+
+def test_dataset_contract_and_roundtrip(tmp_path: Path) -> None:
+    source = _make_source_project(tmp_path / "source")
+    manifest = materialize_subset(source, tmp_path / "target")
+    assert load_manifest(manifest)[0].image_path.is_file()
+    dataset = Pi07SteerableDataset(
+        manifest,
+        image_size=32,
+        prompt_dropout=0.0,
+        max_rtc_delay_steps=12,
+    )
+    assert len(dataset) == 5
+    assert {policy.kind for policy in dataset.policies} == {
+        "expert",
+        "fast",
+        "slow",
+        "understeer",
+        "stop",
+    }
+    item = dataset[0]
+    assert item["pixel_values"].shape == (3, 32, 32)
+    assert item["history_pixel_values"].shape == (6, 3, 32, 32)
+    assert item["history_states"].shape == (6, 2)
+    assert item["current_state"].shape == (2,)
+    assert item["actions"].shape == (16, 2)
+    assert item["subgoal_heatmap"].shape == (1, 32, 32)
+    assert item["subgoal_pixel_values"].shape == (3, 32, 32)
+    assert item["ambiguous_history_pair"].shape == (2, 6, 2)
+    assert 0 <= item["rtc_delay_steps"] <= 12
+    assert item["rtc_target_actions"].shape == (16, 2)
+    assert item["rtc_valid_mask"].shape == (16,)
+    assert item["reasoning"] == "Steer left."
+    assert item["prompt_component_mask"].all()
+    assert item["token_metadata"]["quality"] == 5
+    assert item["control_mode"] == "ee"
+    decoded = dataset.action_normalizer.decode(item["actions"], item["history_xy"])
+    torch.testing.assert_close(decoded, item["future_xy"], atol=1e-5, rtol=1e-5)
+
+
+def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None:
+    source = _make_source_project(tmp_path / "source")
+    dataset = Pi07SteerableDataset(
+        source,
+        image_size=24,
+        prompt_dropout=0.5,
+        seed=11,
+    )
+    first = dataset[3]
+    second = dataset[3]
+    assert first["prompt_components"] == second["prompt_components"]
+    assert first["rtc_delay_steps"] == second["rtc_delay_steps"]
+    torch.testing.assert_close(first["rtc_target_actions"], second["rtc_target_actions"])