seeker
66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
import torch
|
|
|
|
from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer
|
|
|
|
|
|
def metadata() -> EpisodeMetadata:
|
|
return EpisodeMetadata(
|
|
speed_steps=2100,
|
|
quality=5,
|
|
mistake=False,
|
|
control_mode="ee",
|
|
source_type="autonomous_rollout",
|
|
)
|
|
|
|
|
|
def test_prompt_contains_rich_metadata_and_control_tokens() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
prompt, metadata_indices = tokenizer.encode_prompt(
|
|
"drive on the road", "steer_left", metadata()
|
|
)
|
|
assert tokenizer.SPEED_MARKER in prompt
|
|
assert tokenizer.QUALITY_MARKER in prompt
|
|
assert tokenizer.MISTAKE_MARKER in prompt
|
|
assert tokenizer.SOURCE_MARKER in prompt
|
|
assert tokenizer.CONTROL_MARKER in prompt
|
|
assert tokenizer.CONTROL_EE in prompt
|
|
assert len(metadata_indices) == 8
|
|
|
|
|
|
def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
actions = torch.zeros(3, 2)
|
|
conditioned, unconditional = tokenizer.build_metadata_cfg_pair(
|
|
"drive", "keep_lane", actions, metadata()
|
|
)
|
|
assert tokenizer.QUALITY_MARKER in conditioned.input_ids
|
|
assert tokenizer.QUALITY_MARKER not in unconditional.input_ids
|
|
assert tokenizer.CONTROL_EE in conditioned.input_ids
|
|
assert tokenizer.CONTROL_EE in unconditional.input_ids
|
|
assert conditioned.metadata_indices
|
|
assert not unconditional.metadata_indices
|
|
|
|
|
|
def test_flow_prefix_ends_before_fast_stream() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
sequence = tokenizer.build_policy_sequence(
|
|
"drive", "keep_lane", torch.zeros(4, 2), metadata()
|
|
)
|
|
assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
|
|
assert sequence.action_prefix_mask[: sequence.action_start_index].all()
|
|
assert not sequence.action_prefix_mask[sequence.action_start_index :].any()
|
|
assert sequence.labels[sequence.action_start_index - 1] == 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
|
|
|