seeker
82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from pi06mini.experience import (
|
|
build_synthetic_experience,
|
|
discounted_returns,
|
|
n_step_advantages,
|
|
normalize_returns_to_unit_interval,
|
|
quantile_advantage_labels,
|
|
)
|
|
|
|
|
|
def _geometry() -> tuple[torch.Tensor, torch.Tensor]:
|
|
history = torch.stack((torch.linspace(-7.0, 0.0, 8), torch.zeros(8)), dim=-1)
|
|
x = torch.linspace(1.0, 16.0, 16)
|
|
future = torch.stack((x, 0.025 * x.square()), dim=-1)
|
|
return history, future
|
|
|
|
|
|
def test_synthetic_buffer_has_all_experience_sources_and_failure_geometry() -> None:
|
|
history, future = _geometry()
|
|
episodes = build_synthetic_experience(history, future, episode_prefix="pai")
|
|
|
|
assert [episode.source_type for episode in episodes] == [
|
|
"demo_success",
|
|
"autonomous_understeer",
|
|
"autonomous_stop",
|
|
"intervention_correction",
|
|
]
|
|
demo, understeer, stopped, correction = episodes
|
|
torch.testing.assert_close(demo.future_xy, future)
|
|
assert understeer.future_xy[-1, 1] < future[-1, 1]
|
|
assert torch.equal(stopped.future_xy[-1], stopped.future_xy[-2])
|
|
assert correction.intervention_mask.any()
|
|
assert not correction.intervention_mask[0]
|
|
torch.testing.assert_close(correction.future_xy[-1], future[-1])
|
|
assert torch.all(demo.rewards == 0)
|
|
assert float(understeer.rewards.min()) < 0
|
|
assert float(stopped.rewards.min()) < 0
|
|
|
|
|
|
def test_returns_are_normalized_to_closed_negative_unit_interval() -> None:
|
|
rewards = torch.tensor([[0.0, 0.0, 0.0], [-0.2, -0.4, -1.0]])
|
|
returns = discounted_returns(rewards, gamma=1.0)
|
|
normalized = normalize_returns_to_unit_interval(returns)
|
|
|
|
assert normalized.shape == rewards.shape
|
|
assert float(normalized.min()) == -1.0
|
|
assert float(normalized.max()) == 0.0
|
|
assert torch.all(normalized <= 0.0)
|
|
assert torch.all(normalized >= -1.0)
|
|
assert torch.equal(
|
|
normalize_returns_to_unit_interval(torch.ones(3)), torch.zeros(3)
|
|
)
|
|
|
|
|
|
def test_n_step_advantage_bootstraps_and_stops_at_terminal() -> None:
|
|
rewards = torch.tensor([1.0, 2.0, 3.0])
|
|
values = torch.tensor([0.5, 1.0, 1.5, 8.0])
|
|
dones = torch.tensor([False, True, False])
|
|
advantages = n_step_advantages(
|
|
rewards, values, n_step=3, gamma=0.5, dones=dones
|
|
)
|
|
|
|
# t=0 sees r0 + gamma*r1, then the terminal blocks the large bootstrap.
|
|
assert torch.isclose(advantages[0], torch.tensor(1.5))
|
|
assert torch.isclose(advantages[1], torch.tensor(1.0))
|
|
# t=2 reaches the explicit end and therefore has no bootstrap.
|
|
assert torch.isclose(advantages[2], torch.tensor(1.5))
|
|
|
|
|
|
def test_quantile_labels_force_interventions_positive() -> None:
|
|
advantages = torch.tensor([-3.0, -2.0, -1.0, 2.0])
|
|
interventions = torch.tensor([True, False, False, False])
|
|
labels, threshold = quantile_advantage_labels(
|
|
advantages, quantile=0.75, intervention_mask=interventions
|
|
)
|
|
|
|
assert labels.tolist() == [True, False, False, True]
|
|
assert float(threshold) > -1.0
|