Initial pi0.6 minimal validation

tests/test_experience.py

@@ -0,0 +1,81 @@
+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