Align pi0.7 minimal structure with paper

tests/test_context.py

@@ -6,8 +6,8 @@ from pi07mini.context import (
     PolicyMetadata,
     build_ambiguous_history_pair,
     drop_prompt_components,
-    endpoint_subgoal_heatmap,
     format_policy_prompt,
+    future_observation_proxy,
     make_history_proxy_frames,
     prompt_components,
 )
@@ -18,19 +18,13 @@ def _history() -> torch.Tensor:
     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_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: