Initial pi0 minimal validation

2026-07-11 20:04:18 +08:00
commit 4fe0099f3f
18 changed files with 1611 additions and 0 deletions
@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi0mini.data import ActionNormalizer, Pi0OverfitDataset, load_manifest, materialize_subset
+
+
+def _make_source_project(root: Path) -> Path:
+    processed = root / "data" / "processed"
+    image_dir = processed / "images"
+    trajectory_dir = processed / "trajectories"
+    image_dir.mkdir(parents=True)
+    trajectory_dir.mkdir(parents=True)
+
+    pixels = np.zeros((48, 80, 3), dtype=np.uint8)
+    pixels[..., 0] = 64
+    pixels[..., 1] = np.arange(80, dtype=np.uint8)[None, :]
+    Image.fromarray(pixels).save(image_dir / "sample.jpg")
+
+    history_x = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
+    history = np.stack((history_x, 0.02 * history_x, np.zeros(8)), axis=-1)
+    future_x = np.linspace(1.0, 32.0, 16, dtype=np.float32)
+    future = np.stack((future_x, 0.002 * future_x**2, np.zeros(16)), axis=-1)
+    np.savez_compressed(
+        trajectory_dir / "sample.npz", history_xyz=history, future_xyz=future
+    )
+
+    manifest = processed / "overfit_manifest.jsonl"
+    manifest.write_text(
+        json.dumps(
+            {
+                "clip_id": "clip-a",
+                "t0_us": 1234,
+                "image_path": "data/processed/images/sample.jpg",
+                "trajectory_path": "data/processed/trajectories/sample.npz",
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    return manifest
+
+
+def test_action_normalizer_is_reversible() -> None:
+    history = torch.tensor([[-1.0, 0.0], [0.0, 0.1]])
+    future = torch.tensor([[1.0, 0.2], [2.5, 0.4], [4.0, 0.7]])
+    normalizer = ActionNormalizer.fit([history], [future])
+    actions = normalizer.encode(history, future)
+
+    assert actions.shape == (3, 2)
+    assert torch.allclose(normalizer.decode(actions, history), future, atol=1e-6)
+    restored = ActionNormalizer.from_dict(normalizer.as_dict())
+    assert torch.allclose(restored.decode(actions, history), future, atol=1e-6)
+
+
+def test_materialized_dataset_contract_and_determinism(tmp_path: Path) -> None:
+    source_manifest = _make_source_project(tmp_path / "source")
+    copied_manifest = materialize_subset(source_manifest, tmp_path / "target")
+    records = load_manifest(copied_manifest)
+
+    assert len(records) == 1
+    assert records[0].image_path.is_file()
+    assert records[0].trajectory_path.is_file()
+
+    dataset = Pi0OverfitDataset(copied_manifest, seed=19)
+    repeated = Pi0OverfitDataset(
+        copied_manifest,
+        seed=19,
+        action_statistics=dataset.action_normalizer.as_dict(),
+    )
+    assert len(dataset) == 4
+    assert dataset.state_dim == 16
+
+    original = dataset[0]
+    augmented = dataset[1]
+    repeated_augmented = repeated[1]
+    assert set(original) == {
+        "image",
+        "state",
+        "actions",
+        "history_xy",
+        "future_xy",
+        "sample_id",
+    }
+    assert original["image"].shape == (3, 64, 64)
+    assert original["state"].shape == (16,)
+    assert original["actions"].shape == (16, 2)
+    assert original["history_xy"].shape == (8, 2)
+    assert original["future_xy"].shape == (16, 2)
+    assert original["sample_id"].endswith(":aug0")
+    assert augmented["sample_id"].endswith(":aug1")
+    assert not torch.equal(original["future_xy"], augmented["future_xy"])
+    assert torch.equal(augmented["image"], repeated_augmented["image"])
+    assert torch.equal(augmented["state"], repeated_augmented["state"])
+    assert torch.equal(augmented["actions"], repeated_augmented["actions"])
+
+    decoded = dataset.action_normalizer.decode(
+        original["actions"], original["history_xy"]
+    )
+    assert torch.allclose(decoded, original["future_xy"], atol=1e-5)
+
+
+def test_statistics_can_be_saved(tmp_path: Path) -> None:
+    source_manifest = _make_source_project(tmp_path / "source")
+    manifest = materialize_subset(source_manifest, tmp_path / "target")
+    dataset = Pi0OverfitDataset(manifest)
+    path = dataset.save_action_statistics(tmp_path / "stats.json")
+    values = json.loads(path.read_text(encoding="utf-8"))
+
+    assert len(values["mean"]) == 2
+    assert len(values["scale"]) == 2
+    assert min(values["scale"]) >= 0.05
@@ -0,0 +1,65 @@
+import torch
+
+from pi0mini.flow import FlowMatching
+
+
+def _time_view(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
+
+
+def test_training_batch_obeys_linear_path() -> None:
+    flow = FlowMatching()
+    action = torch.tensor([[[1.0, -2.0], [0.5, 3.0]], [[-1.0, 0.0], [2.5, -0.5]]])
+    x_t, t, velocity = flow.training_batch(
+        action, generator=torch.Generator().manual_seed(123)
+    )
+
+    noise = action - velocity
+    expected = (1.0 - _time_view(t, action)) * noise + _time_view(t, action) * action
+    torch.testing.assert_close(x_t, expected)
+    assert x_t.shape == action.shape
+    assert torch.all((t >= 0.0) & (t <= 1.0))
+
+
+def test_training_batch_is_reproducible() -> None:
+    action = torch.randn(3, 16, 2)
+    flow = FlowMatching()
+    a = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
+    b = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
+    for left, right in zip(a, b):
+        torch.testing.assert_close(left, right)
+
+
+def test_training_times_use_low_biased_shifted_beta() -> None:
+    flow = FlowMatching()
+    times = flow.sample_training_times(
+        100_000,
+        device="cpu",
+        dtype=torch.float32,
+        generator=torch.Generator().manual_seed(17),
+    )
+
+    assert torch.all(times >= 0.0)
+    assert torch.all(times <= 0.999)
+    # E[0.999 * (1 - Beta(1.5, 1))] = 0.3996.  This also guards
+    # against accidentally reverting to a uniform or high-time-biased draw.
+    assert abs(float(times.mean()) - 0.3996) < 0.005
+
+
+def test_loss_is_velocity_mse() -> None:
+    flow = FlowMatching()
+    target = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
+    prediction = target + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
+    torch.testing.assert_close(flow.loss(prediction, target), torch.tensor(1.5))
+
+
+def test_exact_vector_field_integrates_to_action() -> None:
+    flow = FlowMatching(inference_steps=10)
+    action = torch.tensor([[[0.25, -0.5], [1.0, 2.0]], [[-2.0, 0.75], [0.0, 0.5]]])
+
+    def exact_field(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        return (action.to(x_t) - x_t) / (1.0 - _time_view(t, x_t))
+
+    sample = flow.sample(exact_field, tuple(action.shape), action.device,
+                         generator=torch.Generator().manual_seed(5))
+    torch.testing.assert_close(sample, action, atol=1e-5, rtol=1e-5)
@@ -0,0 +1,65 @@
+import torch
+
+from pi0mini.model import Pi0Config, Pi0Mini
+
+
+def test_pi0_three_block_attention_mask() -> None:
+    config = Pi0Config(
+        image_size=32,
+        patch_size=16,
+        max_text_tokens=4,
+        max_action_tokens=3,
+        d_model=32,
+        n_heads=4,
+        n_layers=1,
+        time_feature_dim=8,
+    )
+    model = Pi0Mini(config)
+
+    # Two image patches per axis -> 4 image tokens, plus 2 language tokens.
+    prefix_edges, suffix_edges = model._masks(
+        batch=1,
+        prefix_length=6,
+        suffix_length=4,  # one state + three action tokens
+        image_tokens=4,
+        attention_mask=torch.tensor([[1, 1]]),
+        device=torch.device("cpu"),
+    )
+
+    assert prefix_edges.shape == (1, 6, 6)
+    assert torch.all(prefix_edges)
+
+    # Suffix key order is [six prefix keys, state, action0, action1, action2].
+    assert suffix_edges.shape == (1, 4, 10)
+    assert torch.all(suffix_edges[0, 0, :7])
+    assert not torch.any(suffix_edges[0, 0, 7:])
+    assert torch.all(suffix_edges[0, 1:, :])
+
+
+def test_pi0_mask_respects_padded_language_keys() -> None:
+    model = Pi0Mini(
+        Pi0Config(
+            image_size=16,
+            patch_size=16,
+            max_text_tokens=3,
+            max_action_tokens=1,
+            d_model=16,
+            n_heads=4,
+            n_layers=1,
+            time_feature_dim=8,
+        )
+    )
+    prefix_edges, suffix_edges = model._masks(
+        batch=1,
+        prefix_length=4,
+        suffix_length=2,
+        image_tokens=1,
+        attention_mask=torch.tensor([[1, 0, 0]]),
+        device=torch.device("cpu"),
+    )
+
+    # Padded prefix keys are hidden from valid prefix, state, and action rows.
+    assert not prefix_edges[0, 0, 2]
+    assert not prefix_edges[0, 0, 3]
+    assert not suffix_edges[0, :, 2].any()
+    assert not suffix_edges[0, :, 3].any()