Initial pi0.5 minimal validation

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