Initial pi0 minimal validation

tests/test_flow.py

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