Initial pi0.5 minimal validation

tests/test_data.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
+from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary
+
+
+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)
+
+    image_path = image_dir / "sample.jpg"
+    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_path)
+
+    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)
+    trajectory_path = trajectory_dir / "sample.npz"
+    np.savez_compressed(trajectory_path, history_xyz=history, future_xyz=future)
+
+    manifest = processed / "overfit_manifest.jsonl"
+    row = {
+        "clip_id": "clip-a",
+        "t0_us": 1234,
+        "reasoning": "Steer left through the construction lane.",
+        "image_path": "data/processed/images/sample.jpg",
+        "trajectory_path": "data/processed/trajectories/sample.npz",
+    }
+    manifest.write_text(json.dumps(row) + "\n", encoding="utf-8")
+    return manifest
+
+
+def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
+    tokenizer = CoordinateTokenizer(x_range=(-10.0, 10.0), y_range=(-5.0, 5.0))
+    coordinates = torch.tensor([[-10.0, -5.0], [0.0, 0.0], [10.0, 5.0]])
+    tokens = tokenizer.encode(coordinates)
+
+    assert tokens.dtype == torch.long
+    assert tokens.tolist()[0] == [0, 0]
+    assert tokens.tolist()[-1] == [255, 255]
+    assert 0 <= int(tokens.min()) <= int(tokens.max()) <= 255
+    decoded = tokenizer.decode(tokens)
+    assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)
+
+
+def test_subtask_vocabulary_and_prefix_contract() -> None:
+    tokenizer = MiniTokenizer()
+    history = torch.zeros(8, 2)
+    prefix = tokenizer.build_prefix(history)
+
+    assert prefix.shape == (20,)
+    assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE]
+    assert prefix[-1].item() == tokenizer.SUBTASK_QUERY
+    assert prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255
+    target = tokenizer.encode_subtask("Steer left around the cones")
+    assert tokenizer.subtasks.decode(target) == "steer_left"
+    assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"
+
+
+def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path: Path) -> None:
+    source_manifest = _make_source_project(tmp_path / "source")
+    copied_manifest = materialize_subset(source_manifest, tmp_path / "target", max_samples=1)
+    records = load_manifest(copied_manifest)
+
+    assert len(records) == 1
+    assert records[0].image_path.is_file()
+    assert records[0].trajectory_path.is_file()
+
+    dataset = Pi05OverfitDataset(
+        copied_manifest,
+        image_size=(32, 48),
+        augmentation_factor=3,
+        seed=19,
+    )
+    same_dataset = Pi05OverfitDataset(
+        copied_manifest,
+        image_size=(32, 48),
+        augmentation_factor=3,
+        seed=19,
+        action_statistics=dataset.action_normalizer.as_dict(),
+    )
+
+    assert len(dataset) == 3
+    original = dataset[0]
+    augmented = dataset[1]
+    repeated = same_dataset[1]
+    assert original["image"].shape == (3, 32, 48)
+    assert original["prefix_tokens"].shape == (20,)
+    assert original["actions"].shape == (16, 2)
+    assert original["subtask_target"].shape == ()
+    assert original["subtask_label"] == "steer_left"
+    assert original["reasoning"].startswith("Steer left")
+    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["future_xy"], repeated["future_xy"])
+    assert torch.equal(augmented["image"], repeated["image"])
+
+    reconstructed = dataset.action_normalizer.decode(
+        original["actions"], original["history_xy"]
+    )
+    assert torch.allclose(reconstructed, original["future_xy"], atol=1e-5)
+
+    statistics_path = dataset.save_action_statistics(tmp_path / "action_stats.json")
+    statistics = json.loads(statistics_path.read_text(encoding="utf-8"))
+    assert len(statistics["mean"]) == 2
+    assert len(statistics["scale"]) == 2
+    assert min(statistics["scale"]) >= 0.05