seeker
40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
import torch
|
|
|
|
from alpamayo_lite.trajectory_tokens import (
|
|
DiscreteTrajectoryMask,
|
|
TrajectoryVocabulary,
|
|
dequantize_action,
|
|
extract_trajectory_values,
|
|
quantize_action,
|
|
)
|
|
|
|
|
|
def test_quantized_action_round_trip_is_bounded() -> None:
|
|
action = torch.tensor([[-4.0, -1.25], [0.0, 1.5], [3.9, 4.0]])
|
|
tokens = quantize_action(action, size=128, minimum=-4.0, maximum=4.0)
|
|
reconstructed = dequantize_action(tokens, size=128, minimum=-4.0, maximum=4.0)
|
|
assert float((reconstructed - action).abs().max()) <= 8.0 / 127.0
|
|
|
|
|
|
def test_extract_trajectory_values_between_markers() -> None:
|
|
vocabulary = TrajectoryVocabulary(
|
|
size=8,
|
|
start_id=100,
|
|
end_id=107,
|
|
special_token_ids={"<|traj_future_start|>": 200, "<|traj_future_end|>": 201},
|
|
)
|
|
sequence = torch.tensor([1, 200, 103, 107, 100, 201, 2])
|
|
torch.testing.assert_close(
|
|
extract_trajectory_values(sequence, vocabulary), torch.tensor([3, 7, 0])
|
|
)
|
|
|
|
|
|
def test_discrete_trajectory_mask_changes_only_trajectory_range() -> None:
|
|
input_ids = torch.tensor([[1, 2]])
|
|
scores = torch.randn(1, 16)
|
|
original = scores.clone()
|
|
masked = DiscreteTrajectoryMask(start_id=5, size=4)(input_ids, scores)
|
|
assert torch.isneginf(masked[:, 5:9]).all()
|
|
torch.testing.assert_close(masked[:, :5], original[:, :5])
|
|
torch.testing.assert_close(masked[:, 9:], original[:, 9:])
|