Align pi0.7 minimal structure with paper

src/pi07mini/context.py

@@ -21,7 +21,6 @@ class PolicyMetadata:
     speed_steps: int = 1_000
     quality_score: int = 3
     control_mode: str = "ee"
-    source_type: str = "autonomous"
 
     def __post_init__(self) -> None:
         if self.speed_steps <= 0:
@@ -47,7 +46,6 @@ class PolicyMetadata:
             "quality": self.quality_score,
             "mistake": self.mistake != "none",
             "control_mode": self.control_mode,
-            "source_type": self.source_type,
         }
 
 
@@ -106,36 +104,34 @@ def format_policy_prompt(components: Mapping[str, str]) -> str:
     return "\n".join(pieces)
 
 
-def endpoint_subgoal_heatmap(
-    endpoint_xy: torch.Tensor,
-    *,
-    size: int = 32,
-    sigma: float = 1.5,
-    xy_extent: tuple[float, float, float, float] = (-8.0, 64.0, -16.0, 16.0),
+def future_observation_proxy(
+    image: torch.Tensor,
+    history_xy: torch.Tensor,
+    future_xy: torch.Tensor,
 ) -> torch.Tensor:
-    """Render a normalized Gaussian subgoal heatmap from one endpoint only.
+    """Create a tiny future-camera proxy instead of encoding coordinates as a heatmap.
 
-    It deliberately accepts no trajectory tensor. This API makes it difficult
-    to accidentally leak the entire ground-truth action chunk into the visual
-    subgoal condition.
+    The source dataset contains only the current frame.  For the architecture
+    validation we therefore warp that frame according to the segment endpoint,
+    producing an image-space target with the same contract as ``o_tend`` in the
+    paper.  The policy only receives the generated image, never the coordinates.
     """
 
-    endpoint = torch.as_tensor(endpoint_xy, dtype=torch.float32).reshape(2)
-    if size < 2 or sigma <= 0.0:
-        raise ValueError("size must be >=2 and sigma must be positive")
-    x_min, x_max, y_min, y_max = xy_extent
-    if not x_min < x_max or not y_min < y_max:
-        raise ValueError("xy_extent must contain increasing bounds")
-    px = (endpoint[0] - x_min) / (x_max - x_min) * (size - 1)
-    py = (y_max - endpoint[1]) / (y_max - y_min) * (size - 1)
-    yy, xx = torch.meshgrid(
-        torch.arange(size, dtype=torch.float32),
-        torch.arange(size, dtype=torch.float32),
-        indexing="ij",
-    )
-    heatmap = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2.0 * sigma**2))
-    heatmap = heatmap / heatmap.amax().clamp_min(torch.finfo(heatmap.dtype).eps)
-    return heatmap.unsqueeze(0)
+    image = torch.as_tensor(image, dtype=torch.float32)
+    history = torch.as_tensor(history_xy, dtype=torch.float32)
+    future = torch.as_tensor(future_xy, dtype=torch.float32)
+    if image.ndim != 3 or image.shape[0] not in {1, 3}:
+        raise ValueError("image must have shape [C,H,W]")
+    if history.ndim != 2 or future.ndim != 2 or history.shape[-1] != 2 or future.shape[-1] != 2:
+        raise ValueError("history_xy and future_xy must have shape [T,2]")
+    displacement = future[-1] - history[-1]
+    scale = max(image.shape[-2:]) / 32.0
+    dx = int(torch.round(displacement[1] * scale).clamp(-image.shape[-1] // 5, image.shape[-1] // 5))
+    dy = int(torch.round(-displacement[0] * scale * 0.08).clamp(-image.shape[-2] // 6, image.shape[-2] // 6))
+    distance = torch.linalg.vector_norm(displacement)
+    exposure = float((1.0 + 0.02 * torch.tanh(distance / 8.0)).item())
+    shifted = torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure
+    return F.avg_pool2d(shifted[None], kernel_size=3, stride=1, padding=1)[0].clamp(0.0, 1.0)
 
 
 def make_history_proxy_frames(