Initial pi0.7 minimal validation

tests/test_context.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import torch
+
+from pi07mini.context import (
+    PolicyMetadata,
+    build_ambiguous_history_pair,
+    drop_prompt_components,
+    endpoint_subgoal_heatmap,
+    format_policy_prompt,
+    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_subgoal_heatmap_depends_on_endpoint_only() -> None:
+    endpoint = torch.tensor([20.0, 4.0])
+    first = endpoint_subgoal_heatmap(endpoint, size=24)
+    # Two future paths with radically different interiors produce exactly the
+    # same condition because only their shared endpoint is passed to the API.
+    path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
+    path_b = path_a.clone()
+    path_b[1:-1, 1] += 100.0
+    second = endpoint_subgoal_heatmap(path_b[-1], size=24)
+    torch.testing.assert_close(first, second, atol=0.0, rtol=0.0)
+    assert first.shape == (1, 24, 24)
+    assert float(first.max()) == 1.0
+    assert torch.isfinite(first).all()
+
+
+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]