Initial pi0 minimal validation

tests/test_data.py

@@ -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