Initial pi0.6 minimal validation

tests/test_value.py

@@ -0,0 +1,73 @@
+import torch
+
+from pi06mini.model import Pi06Config, Pi06Mini
+from pi06mini.tokenizer import MiniTokenizer
+from pi06mini.value import DistributionalValueModel, ValueConfig
+
+
+def tiny_value_model() -> DistributionalValueModel:
+    tokenizer = MiniTokenizer()
+    return DistributionalValueModel(
+        ValueConfig(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            max_text_tokens=16,
+        )
+    )
+
+
+def test_distributional_value_has_201_bins_and_expected_value() -> None:
+    tokenizer = MiniTokenizer()
+    model = tiny_value_model()
+    state = tokenizer.build_value_input(torch.zeros(2, 2))[None].repeat(2, 1)
+    output = model(
+        torch.randn(2, 3, 16, 16),
+        state,
+        target_values=torch.tensor([-1.0, 0.0]),
+    )
+
+    assert output.logits.shape == (2, 201)
+    assert output.probabilities.shape == (2, 201)
+    torch.testing.assert_close(output.probabilities.sum(-1), torch.ones(2))
+    assert torch.all((-1.0 <= output.expectation) & (output.expectation <= 0.0))
+    assert output.target_bins.tolist() == [0, 200]
+    assert output.loss is not None and output.loss.ndim == 0
+
+
+def test_value_discretization_and_expectation_use_configured_support() -> None:
+    model = tiny_value_model()
+    bins = model.values_to_bins(torch.tensor([-2.0, -0.5, 1.0]))
+    assert bins.tolist() == [0, 100, 200]
+    torch.testing.assert_close(
+        model.bins_to_values(torch.tensor([0, 100, 200])),
+        torch.tensor([-1.0, -0.5, 0.0]),
+    )
+    logits = torch.full((3, 201), -100.0)
+    logits[0, 0] = logits[1, 100] = logits[2, 200] = 100.0
+    torch.testing.assert_close(
+        model.expectation_from_logits(logits),
+        torch.tensor([-1.0, -0.5, 0.0]),
+    )
+
+
+def test_value_model_and_policy_share_no_parameters() -> None:
+    tokenizer = MiniTokenizer()
+    policy = Pi06Mini(
+        Pi06Config(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            action_start_token_id=tokenizer.ACTION_START,
+        )
+    )
+    value = tiny_value_model()
+    policy_ids = {id(parameter) for parameter in policy.parameters()}
+    value_ids = {id(parameter) for parameter in value.parameters()}
+    assert policy_ids.isdisjoint(value_ids)