seeker
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
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,
|
|
]
|
|
)
|
|
)
|