seeker
51 lines
1.8 KiB
Python
51 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay
|
|
|
|
|
|
def _actions() -> torch.Tensor:
|
|
return torch.arange(32, dtype=torch.float32).reshape(16, 2)
|
|
|
|
|
|
def test_zero_delay_preserves_entire_chunk() -> None:
|
|
actions = _actions()
|
|
aligned = align_actions_for_rtc(actions, 0)
|
|
torch.testing.assert_close(aligned.delayed_target_actions, actions)
|
|
assert aligned.valid_mask.all()
|
|
assert aligned.executed_prefix_actions.shape == (0, 2)
|
|
|
|
|
|
def test_delay_shifts_target_and_masks_padded_tail() -> None:
|
|
actions = _actions()
|
|
aligned = align_actions_for_rtc(actions, 12)
|
|
torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:])
|
|
assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0
|
|
assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12
|
|
torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12])
|
|
|
|
|
|
def test_rtc_alignment_preserves_batch_dimensions() -> None:
|
|
actions = torch.stack((_actions(), _actions() + 100.0))
|
|
aligned = align_actions_for_rtc(actions, 3)
|
|
assert aligned.delayed_target_actions.shape == (2, 16, 2)
|
|
assert aligned.valid_mask.shape == (2, 16)
|
|
torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:])
|
|
|
|
|
|
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):
|
|
align_actions_for_rtc(_actions(), delay)
|