seeker
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from pi07mini.flow import FlowMatching
|
|
from pi07mini.rtc import build_rtc_conditioning, sample_rtc_delay
|
|
|
|
|
|
def _actions() -> torch.Tensor:
|
|
return torch.arange(32, dtype=torch.float32).reshape(16, 2)
|
|
|
|
|
|
def test_zero_delay_has_empty_prefix_and_full_postfix() -> None:
|
|
rtc = build_rtc_conditioning(_actions(), 0)
|
|
assert not rtc.prefix_mask.any()
|
|
assert rtc.postfix_mask.all()
|
|
assert torch.count_nonzero(rtc.action_prefix) == 0
|
|
|
|
|
|
def test_delay_keeps_ground_truth_prefix_without_shifting_target() -> None:
|
|
actions = _actions()
|
|
rtc = build_rtc_conditioning(actions, 12)
|
|
torch.testing.assert_close(rtc.action_prefix[:12], actions[:12])
|
|
assert torch.count_nonzero(rtc.action_prefix[12:]) == 0
|
|
assert rtc.prefix_mask.tolist() == [True] * 12 + [False] * 4
|
|
assert rtc.postfix_mask.tolist() == [False] * 12 + [True] * 4
|
|
|
|
|
|
def test_flow_training_sets_prefix_time_to_one_and_masks_its_loss() -> None:
|
|
actions = _actions()[None]
|
|
noisy, times, _, loss_mask = FlowMatching().training_batch(
|
|
actions, delays=torch.tensor([3]), generator=torch.Generator().manual_seed(4)
|
|
)
|
|
torch.testing.assert_close(noisy[:, :3], actions[:, :3])
|
|
torch.testing.assert_close(times[:, :3], torch.ones(1, 3))
|
|
assert not loss_mask[:, :3].any()
|
|
assert loss_mask[:, 3:].all()
|
|
|
|
|
|
def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
|
|
generator = torch.Generator().manual_seed(9)
|
|
values = {
|
|
sample_rtc_delay(generator=generator, action_horizon=16) for _ in range(300)
|
|
}
|
|
assert min(values) == 0
|
|
assert max(values) == 12
|
|
|
|
|
|
@pytest.mark.parametrize("delay", [-1, 13, 16])
|
|
def test_invalid_delay_is_rejected(delay: int) -> None:
|
|
with pytest.raises(ValueError):
|
|
build_rtc_conditioning(_actions(), delay)
|