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,
+                ]
+            )
+        )