Initial pi0.6 minimal validation

tests/test_model.py

@@ -0,0 +1,123 @@
+import torch
+
+from pi06mini.model import Pi06Config, Pi06Mini
+from pi06mini.tokenizer import MiniTokenizer
+
+
+def tiny_model() -> Pi06Mini:
+    tokenizer = MiniTokenizer()
+    return Pi06Mini(
+        Pi06Config(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            max_text_tokens=32,
+            max_action_tokens=4,
+            action_start_token_id=tokenizer.ACTION_START,
+        )
+    )
+
+
+def test_flow_reads_subtask_and_advantage_but_not_fast_targets() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
+    )
+    model = tiny_model()
+    input_ids = sequence.input_ids[None]
+    derived = model._derive_action_prefix_mask(input_ids)
+    torch.testing.assert_close(derived[0], sequence.action_prefix_mask)
+
+    text_offset = 1
+    prefix, flow = model._masks(
+        batch=1,
+        prefix_length=1 + input_ids.shape[1],
+        action_length=2,
+        text_offset=text_offset,
+        attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
+        action_prefix_mask=derived,
+        generated_token_start=sequence.generated_token_start,
+        device=torch.device("cpu"),
+    )
+    generated_start = text_offset + sequence.generated_token_start
+    assert prefix[0, 0, :generated_start].all()
+    assert not prefix[0, 0, generated_start:].any()
+    # Flow sees image plus all text strictly before ACTION_START.
+    readable_prefix = text_offset + sequence.action_start_index
+    assert flow[0, :, :readable_prefix].all()
+    assert not flow[0, :, readable_prefix : 1 + input_ids.shape[1]].any()
+    assert flow[0, :, 1 + input_ids.shape[1] :].all()
+
+
+def test_flow_loss_is_knowledge_insulated_from_every_vlm_parameter() -> None:
+    torch.manual_seed(4)
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
+    )
+    model = tiny_model()
+    output = model(
+        pixel_values=torch.randn(1, 3, 16, 16),
+        input_ids=sequence.input_ids[None],
+        noisy_actions=torch.randn(1, 2, 2),
+        timesteps=torch.rand(1),
+        action_prefix_mask=sequence.action_prefix_mask[None],
+        generated_token_start=sequence.generated_token_start,
+        target_velocity=torch.randn(1, 2, 2),
+    )
+    assert output.flow_loss is not None
+    output.flow_loss.backward()
+
+    vlm_gradients = [parameter.grad for parameter in model.vlm_parameters()]
+    assert vlm_gradients
+    assert all(gradient is None or torch.count_nonzero(gradient) == 0 for gradient in vlm_gradients)
+    action_gradients = [parameter.grad for parameter in model.action_parameters()]
+    assert any(gradient is not None and torch.count_nonzero(gradient) > 0 for gradient in action_gradients)
+
+
+def test_knowledge_insulation_stops_gradients_not_conditioning() -> None:
+    torch.manual_seed(8)
+    tokenizer = MiniTokenizer()
+    model = tiny_model().eval()
+    positive = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "positive"
+    )
+    negative = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "negative"
+    )
+    image = torch.randn(1, 3, 16, 16)
+    noise = torch.randn(1, 2, 2)
+    timestep = torch.full((1,), 0.5)
+    with torch.no_grad():
+        pos = model(
+            image, positive.input_ids[None], noise, timestep,
+            generated_token_start=positive.generated_token_start,
+        ).vector_field
+        neg = model(
+            image, negative.input_ids[None], noise, timestep,
+            generated_token_start=negative.generated_token_start,
+        ).vector_field
+    assert not torch.allclose(pos, neg)
+
+
+def test_joint_loss_still_trains_vlm_through_token_objective() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "stop", torch.zeros(2, 2), "negative"
+    )
+    model = tiny_model()
+    output = model(
+        torch.randn(1, 3, 16, 16),
+        sequence.input_ids[None],
+        torch.randn(1, 2, 2),
+        torch.rand(1),
+        generated_token_start=sequence.generated_token_start,
+        text_labels=sequence.labels[None],
+        target_velocity=torch.randn(1, 2, 2),
+    )
+    output.loss.backward()
+    assert model.token_embedding.weight.grad is not None
+    assert torch.count_nonzero(model.token_embedding.weight.grad) > 0