seeker
71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from pi07mini.context import (
|
|
PolicyMetadata,
|
|
build_ambiguous_history_pair,
|
|
drop_prompt_components,
|
|
format_policy_prompt,
|
|
future_observation_proxy,
|
|
make_history_proxy_frames,
|
|
prompt_components,
|
|
)
|
|
|
|
|
|
def _history() -> torch.Tensor:
|
|
x = torch.linspace(-8.0, 0.0, 8)
|
|
return torch.stack((x, 0.04 * x.square()), -1)
|
|
|
|
|
|
def test_future_subgoal_has_image_contract_not_coordinate_heatmap() -> None:
|
|
image = torch.linspace(0.0, 1.0, 3 * 24 * 24).reshape(3, 24, 24)
|
|
future = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
|
|
subgoal = future_observation_proxy(image, _history(), future)
|
|
assert subgoal.shape == image.shape
|
|
assert torch.isfinite(subgoal).all()
|
|
assert not torch.equal(subgoal, image)
|
|
|
|
|
|
def test_history_proxy_is_oldest_to_newest_and_future_free() -> None:
|
|
image = torch.linspace(0.0, 1.0, 3 * 20 * 20).reshape(3, 20, 20)
|
|
frames = make_history_proxy_frames(image, _history(), frames=6)
|
|
assert frames.shape == (6, 3, 20, 20)
|
|
torch.testing.assert_close(frames[-1], image * 0.98)
|
|
assert not torch.equal(frames[0], frames[-1])
|
|
|
|
|
|
def test_ambiguous_histories_are_distinct_but_share_current_state() -> None:
|
|
pair = build_ambiguous_history_pair(_history(), frames=6)
|
|
assert pair.shape == (2, 6, 2)
|
|
torch.testing.assert_close(pair[0, -1], pair[1, -1])
|
|
assert not torch.allclose(pair[0, :-1], pair[1, :-1])
|
|
|
|
|
|
def test_prompt_component_dropout_is_independent_and_deterministic() -> None:
|
|
metadata = PolicyMetadata("fast", "good", "none", "autonomous")
|
|
components = prompt_components(metadata, task="Drive safely.")
|
|
generator_a = torch.Generator().manual_seed(42)
|
|
generator_b = torch.Generator().manual_seed(42)
|
|
visible_a, mask_a = drop_prompt_components(
|
|
components, drop_probability=0.5, generator=generator_a
|
|
)
|
|
visible_b, mask_b = drop_prompt_components(
|
|
components, drop_probability=0.5, generator=generator_b
|
|
)
|
|
assert visible_a == visible_b
|
|
torch.testing.assert_close(mask_a, mask_b)
|
|
assert visible_a["task"] == "Drive safely."
|
|
assert mask_a.shape == (5,)
|
|
text = format_policy_prompt(visible_a)
|
|
assert "Drive safely." in text
|
|
for dropped in set(components) - set(visible_a):
|
|
assert f"{dropped}=" not in text
|
|
|
|
|
|
def test_full_metadata_dropout_keeps_only_task() -> None:
|
|
components = prompt_components(PolicyMetadata("slow", "poor", "stop", "autonomous"))
|
|
visible, mask = drop_prompt_components(components, drop_probability=1.0)
|
|
assert set(visible) == {"task"}
|
|
assert mask.tolist() == [True, False, False, False, False]
|