seeker
66 lines
2.7 KiB
Python
66 lines
2.7 KiB
Python
import torch
|
|
|
|
from pi06mini.tokenizer import MiniTokenizer
|
|
|
|
|
|
def test_advantage_condition_is_between_subtask_and_actions() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
history = torch.tensor([[0.0, 0.0], [1.0, 0.2]])
|
|
actions = torch.tensor([[0.1, 0.0], [0.2, -0.1]])
|
|
|
|
positive = tokenizer.build_policy_sequence(
|
|
history, "steer_left", actions, "positive"
|
|
)
|
|
assert positive.advantage_index == positive.subtask_index + 1
|
|
assert positive.action_start_index == positive.advantage_index + 1
|
|
assert positive.input_ids[positive.advantage_index] == tokenizer.ADVANTAGE_POSITIVE
|
|
assert positive.input_ids[positive.action_start_index] == tokenizer.ACTION_START
|
|
assert positive.action_prefix_mask[: positive.action_start_index].all()
|
|
assert not positive.action_prefix_mask[positive.action_start_index :].any()
|
|
|
|
negative = tokenizer.build_policy_sequence(
|
|
history, "steer_left", actions, "negative"
|
|
)
|
|
assert negative.input_ids[negative.advantage_index] == tokenizer.ADVANTAGE_NEGATIVE
|
|
|
|
|
|
def test_omitted_advantage_is_absent_not_a_third_condition_token() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
sequence = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "omitted"
|
|
)
|
|
|
|
assert sequence.advantage_index is None
|
|
assert sequence.action_start_index == sequence.subtask_index + 1
|
|
assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
|
|
assert tokenizer.ADVANTAGE_POSITIVE not in sequence.input_ids
|
|
assert tokenizer.ADVANTAGE_NEGATIVE not in sequence.input_ids
|
|
|
|
|
|
def test_labels_predict_subtask_and_fast_but_not_external_advantage() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
positive = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "stop", torch.zeros(2, 2), True
|
|
)
|
|
assert positive.labels[positive.subtask_index - 1] == positive.input_ids[positive.subtask_index]
|
|
assert positive.labels[positive.subtask_index] == -100
|
|
assert positive.labels[positive.advantage_index] == tokenizer.ACTION_START
|
|
|
|
omitted = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "stop", torch.zeros(2, 2), None
|
|
)
|
|
assert omitted.labels[omitted.subtask_index] == tokenizer.ACTION_START
|
|
|
|
|
|
def test_fast_lite_dct_round_trip_has_bounded_error() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
actions = torch.tensor(
|
|
[[0.0, 0.0], [0.2, -0.1], [0.5, -0.2], [0.8, -0.1]],
|
|
dtype=torch.float32,
|
|
)
|
|
exact = tokenizer.idct_actions(tokenizer.dct_actions(actions))
|
|
torch.testing.assert_close(exact, actions, atol=1e-5, rtol=1e-5)
|
|
|
|
decoded = tokenizer.decode_actions(tokenizer.encode_actions(actions))
|
|
assert torch.sqrt(torch.mean((decoded - actions) ** 2)) < 0.3
|