Initial pi0.5 minimal validation

tests/test_action_tokens.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from pi05mini.tokenizer import MiniTokenizer
+
+
+def test_orthonormal_dct_round_trip() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.tensor(
+        [
+            [-4.0, -2.75],
+            [-0.25, 0.0],
+            [1.375, 3.9],
+            [4.0, 2.125],
+        ]
+    )
+
+    coefficients = tokenizer.dct_actions(actions)
+    reconstructed = tokenizer.idct_actions(coefficients)
+
+    assert coefficients.shape == actions.shape
+    torch.testing.assert_close(reconstructed, actions, atol=1e-5, rtol=1e-5)
+
+
+def test_fast_lite_tokens_round_trip_with_bounded_quantization_error() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.tensor(
+        [
+            [-3.5, -2.75],
+            [-0.25, 0.0],
+            [1.375, 3.5],
+            [3.75, 2.125],
+        ]
+    )
+
+    tokens = tokenizer.encode_actions(actions)
+    reconstructed = tokenizer.decode_actions(tokens)
+
+    bin_width = (
+        tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM
+    ) / (tokenizer.ACTION_BINS - 1)
+    assert reconstructed.shape == actions.shape
+    # Orthonormal IDCT preserves total coefficient error.  Coefficients are
+    # quantized after division by sqrt(horizon), yielding this RMSE bound.
+    rmse = torch.sqrt(torch.mean((reconstructed - actions).square()))
+    assert float(rmse) <= bin_width / 2 * actions.shape[0] ** 0.5 + 1e-6
+
+
+def test_action_tokens_are_serialized_low_frequency_first() -> None:
+    tokenizer = MiniTokenizer()
+    horizon = 4
+    coefficients = torch.zeros(horizon, 2)
+    coefficients[0, 0] = 1.5 * horizon**0.5  # positive X DC coefficient
+    coefficients[1, 1] = -1.5 * horizon**0.5  # negative Y first harmonic
+    actions = tokenizer.idct_actions(coefficients)
+
+    payload = tokenizer.encode_actions(actions)[1:-1]
+    zero_token = tokenizer.ACTION_TOKEN_START + round(
+        (0.0 - tokenizer.ACTION_MINIMUM)
+        / (tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM)
+        * (tokenizer.ACTION_BINS - 1)
+    )
+
+    # Frequency-major XY order is [DC-x, DC-y, f1-x, f1-y, ...].
+    assert payload[0].item() > zero_token
+    assert payload[1].item() == zero_token
+    assert payload[2].item() == zero_token
+    assert payload[3].item() < zero_token
+    assert torch.all(payload[4:] == zero_token)
+
+
+def test_action_token_ids_are_contiguous_reserved_vocabulary() -> None:
+    tokenizer = MiniTokenizer()
+    edge_actions = torch.tensor([[-4.0, 4.0], [0.0, 1.0]])
+    tokens = tokenizer.encode_actions(edge_actions)
+
+    assert tokens[0].item() == tokenizer.ACTION_START
+    assert tokens[-1].item() == tokenizer.ACTION_END
+    assert tokens[1:-1].min().item() >= tokenizer.ACTION_TOKEN_START
+    assert tokens[1:-1].max().item() <= tokenizer.ACTION_TOKEN_END
+    assert tuple(tokenizer.action_token_ids) == tuple(
+        range(tokenizer.ACTION_TOKEN_START, tokenizer.ACTION_TOKEN_END + 1)
+    )
+    assert tokenizer.vocab_size == tokenizer.ACTION_TOKEN_END + 1
+
+
+def test_decode_actions_rejects_missing_boundaries_and_non_action_ids() -> None:
+    tokenizer = MiniTokenizer()
+
+    with pytest.raises(ValueError, match="enclosed"):
+        tokenizer.decode_actions(torch.tensor([tokenizer.ACTION_TOKEN_START] * 2))
+    with pytest.raises(ValueError, match="reserved"):
+        tokenizer.decode_actions(
+            torch.tensor(
+                [
+                    tokenizer.ACTION_START,
+                    tokenizer.BOS,
+                    tokenizer.ACTION_TOKEN_START,
+                    tokenizer.ACTION_END,
+                ]
+            )
+        )

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

tests/test_flow.py

@@ -0,0 +1,85 @@
+import torch
+
+from pi05mini.flow import FlowMatching
+
+
+def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Expand a per-example time tensor over non-batch dimensions."""
+    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
+
+
+def test_training_batch_obeys_linear_flow_equations() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[1.0, -2.0], [0.5, 3.0]],
+            [[-1.0, 0.0], [2.5, -0.5]],
+        ]
+    )
+    generator = torch.Generator().manual_seed(123)
+
+    x_t, t, target_velocity = flow.training_batch(target, generator=generator)
+
+    assert x_t.shape == target.shape
+    assert target_velocity.shape == target.shape
+    assert t.shape[0] == target.shape[0]
+    assert torch.all((t >= 0.0) & (t <= 1.0))
+
+    t_view = _expand_batch_time(t, target)
+    # v = target - noise, so the sampled noise is exactly recoverable.
+    noise = target - target_velocity
+    expected_x_t = (1.0 - t_view) * noise + t_view * target
+    torch.testing.assert_close(x_t, expected_x_t)
+
+
+def test_training_batch_is_reproducible_with_generator() -> None:
+    flow = FlowMatching()
+    target = torch.randn(3, 5, 2)
+
+    batch_a = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+    batch_b = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+
+    for tensor_a, tensor_b in zip(batch_a, batch_b):
+        torch.testing.assert_close(tensor_a, tensor_b)
+
+
+def test_loss_is_mean_squared_velocity_error() -> None:
+    flow = FlowMatching()
+    target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
+    prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
+
+    # Squared errors are [1, 1, 4, 0], whose mean is 1.5.
+    torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5))
+    torch.testing.assert_close(
+        flow.loss(target_velocity, target_velocity), torch.tensor(0.0)
+    )
+
+
+def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[0.25, -0.5], [1.0, 2.0]],
+            [[-2.0, 0.75], [0.0, 0.5]],
+        ]
+    )
+
+    def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        t_view = _expand_batch_time(t, x_t)
+        # Straight conditional path to target. The sampler evaluates t < 1.
+        return (target.to(x_t) - x_t) / (1.0 - t_view)
+
+    torch.manual_seed(5)
+    sampled = flow.sample(
+        exact_model,
+        shape=tuple(target.shape),
+        device=target.device,
+        steps=8,
+    )
+
+    torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5)
+

tests/test_model.py

@@ -0,0 +1,61 @@
+import torch
+
+from pi05mini.model import Pi05Config, Pi05Mini
+
+
+def tiny_model() -> Pi05Mini:
+    return Pi05Mini(Pi05Config(
+        vocab_size=300,
+        image_size=16,
+        patch_size=16,
+        d_model=32,
+        n_heads=4,
+        n_layers=1,
+        max_text_tokens=16,
+        max_action_tokens=4,
+    ))
+
+
+def test_pi05_has_no_continuous_state_projection() -> None:
+    model = tiny_model()
+    assert not hasattr(model, "state_projection")
+    output = model(
+        torch.randn(2, 3, 16, 16),
+        torch.randint(0, 300, (2, 7)),
+        torch.randn(2, 4, 2),
+        torch.rand(2),
+        generated_token_start=4,
+    )
+    assert output.text_logits.shape == (2, 7, 300)
+    assert output.vector_field.shape == (2, 4, 2)
+
+
+def test_observation_generated_and_flow_attention_partition() -> None:
+    model = tiny_model()
+    # One image patch followed by four observation tokens and three generated
+    # tokens. Flow may read observation plus the first generated (subtask)
+    # token, but not the two teacher-forced FAST tokens.
+    action_prefix_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], dtype=torch.bool)
+    prefix, action = model._masks(
+        batch=1,
+        prefix_length=8,
+        action_length=2,
+        text_offset=1,
+        attention_mask=torch.ones(1, 7, dtype=torch.bool),
+        action_prefix_mask=action_prefix_mask,
+        generated_token_start=4,
+        device=torch.device("cpu"),
+    )
+
+    generated_start = 5  # image patch + four observation text tokens
+    # Observation block is fully bidirectional and cannot see generated tokens.
+    assert prefix[0, 0, :generated_start].all()
+    assert not prefix[0, 0, generated_start:].any()
+    assert prefix[0, 4, :generated_start].all()
+    # Generated tokens see all observation tokens and only prior/self generated tokens.
+    assert prefix[0, 5, :6].all() and not prefix[0, 5, 6:].any()
+    assert prefix[0, 6, :7].all() and not prefix[0, 6, 7]
+    # Flow reads allowed prefix keys and all flow-action keys bidirectionally.
+    assert action[0, :, :6].all()
+    assert not action[0, :, 6:8].any()
+    assert action[0, :, 8:].all()