Initial pi0.7 minimal validation

2026-07-12 00:52:00 +08:00
commit b4c8db846b
27 changed files with 3724 additions and 0 deletions
@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Mapping, Sequence
+
+import torch
+import torch.nn.functional as F
+
+
+PROMPT_COMPONENT_ORDER = ("task", "speed", "quality", "mistake", "control")
+
+
+@dataclass(frozen=True)
+class PolicyMetadata:
+    """Steerable policy attributes used by the tiny pi0.7 experiment."""
+
+    speed: str
+    quality: str
+    mistake: str
+    control: str
+    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:
+            raise ValueError("speed_steps must be positive")
+        if not 1 <= self.quality_score <= 5:
+            raise ValueError("quality_score must lie in [1,5]")
+        if self.control_mode not in {"joint", "ee"}:
+            raise ValueError("control_mode must be joint or ee")
+
+    def as_dict(self) -> dict[str, str]:
+        return {
+            "speed": self.speed,
+            "quality": self.quality,
+            "mistake": self.mistake,
+            "control": self.control,
+        }
+
+    def as_token_metadata(self) -> dict[str, int | bool | str]:
+        """Return the scalar metadata contract consumed by MiniTokenizer."""
+
+        return {
+            "speed_steps": self.speed_steps,
+            "quality": self.quality_score,
+            "mistake": self.mistake != "none",
+            "control_mode": self.control_mode,
+            "source_type": self.source_type,
+        }
+
+
+def prompt_components(
+    metadata: PolicyMetadata,
+    *,
+    task: str = "Follow the visible road and predict the next waypoint chunk.",
+) -> dict[str, str]:
+    return {"task": task, **metadata.as_dict()}
+
+
+def drop_prompt_components(
+    components: Mapping[str, str],
+    *,
+    drop_probability: float,
+    generator: torch.Generator | None = None,
+    preserve: Sequence[str] = ("task",),
+) -> tuple[dict[str, str], torch.Tensor]:
+    """Apply classifier-free-style dropout to independent prompt attributes.
+
+    The returned mask follows ``PROMPT_COMPONENT_ORDER`` and marks components
+    that remain visible. Task text is preserved by default, while metadata can
+    be independently omitted so the policy cannot require every attribute.
+    """
+
+    if not 0.0 <= drop_probability <= 1.0:
+        raise ValueError("drop_probability must lie in [0, 1]")
+    preserved = set(preserve)
+    visible: dict[str, str] = {}
+    mask = torch.zeros(len(PROMPT_COMPONENT_ORDER), dtype=torch.bool)
+    for index, name in enumerate(PROMPT_COMPONENT_ORDER):
+        if name not in components:
+            continue
+        keep = name in preserved or bool(
+            torch.rand((), generator=generator) >= drop_probability
+        )
+        if keep:
+            visible[name] = str(components[name])
+            mask[index] = True
+    return visible, mask
+
+
+def format_policy_prompt(components: Mapping[str, str]) -> str:
+    """Format components without inventing values for dropped attributes."""
+
+    pieces: list[str] = []
+    if task := components.get("task"):
+        pieces.append(str(task).strip())
+    attributes = [
+        f"{name}={components[name]}"
+        for name in PROMPT_COMPONENT_ORDER[1:]
+        if name in components
+    ]
+    if attributes:
+        pieces.append("Policy metadata: " + ", ".join(attributes) + ".")
+    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),
+) -> torch.Tensor:
+    """Render a normalized Gaussian subgoal heatmap from one endpoint only.
+
+    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.
+    """
+
+    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)
+
+
+def make_history_proxy_frames(
+    image: torch.Tensor,
+    history_xy: torch.Tensor,
+    *,
+    frames: int = 6,
+) -> torch.Tensor:
+    """Create oldest-to-newest history frames from one PAI image and past XY.
+
+    PAI's tiny overfit manifest only exposes one camera frame. We therefore
+    synthesize a deterministic visual-history proxy by rolling that frame from
+    past ego displacement and applying a very small temporal exposure change.
+    No future waypoint is accepted by this function.
+    """
+
+    image = torch.as_tensor(image, dtype=torch.float32)
+    history = torch.as_tensor(history_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 history.shape[-1] != 2 or len(history) < 2:
+        raise ValueError("history_xy must have shape [T>=2,2]")
+    if frames < 2:
+        raise ValueError("frames must be >=2")
+    sampled = F.interpolate(
+        history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True
+    )[0].T
+    relative = sampled - sampled[-1]
+    spatial_scale = max(image.shape[-2:]) / 24.0
+    output: list[torch.Tensor] = []
+    for index, point in enumerate(relative):
+        dx = int(torch.round(point[1] * spatial_scale).item())
+        dy = int(torch.round(-point[0] * spatial_scale * 0.25).item())
+        dx = max(-image.shape[-1] // 6, min(image.shape[-1] // 6, dx))
+        dy = max(-image.shape[-2] // 6, min(image.shape[-2] // 6, dy))
+        # The separate current image is appended by MEM, so even the newest
+        # proxy frame remains a slightly older exposure rather than an exact
+        # duplicate of current_pixel_values.
+        exposure = 0.90 + 0.08 * index / (frames - 1)
+        output.append(torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure)
+    return torch.stack(output).clamp(0.0, 1.0)
+
+
+def build_ambiguous_history_pair(
+    history_xy: torch.Tensor, *, frames: int = 6
+) -> torch.Tensor:
+    """Return two distinct histories that share exactly the same current state.
+
+    This supplies a minimal history-disambiguation probe: a current-frame-only
+    policy sees the same endpoint, whereas a memory encoder can distinguish the
+    two past paths. Construction depends only on observed history.
+    """
+
+    history = torch.as_tensor(history_xy, dtype=torch.float32)
+    if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2:
+        raise ValueError("history_xy must have shape [T>=2,2]")
+    sampled = F.interpolate(
+        history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True
+    )[0].T
+    alpha = torch.linspace(0.0, 1.0, frames, dtype=sampled.dtype).unsqueeze(-1)
+    chord = sampled[:1] + alpha * (sampled[-1:] - sampled[:1])
+    mirrored = 2.0 * chord - sampled
+    if torch.allclose(mirrored, sampled, atol=1e-4, rtol=0.0):
+        direction = sampled[-1] - sampled[0]
+        perpendicular = torch.stack((-direction[1], direction[0]))
+        perpendicular = perpendicular / torch.linalg.vector_norm(perpendicular).clamp_min(1e-6)
+        amplitude = torch.linalg.vector_norm(direction).clamp_min(1.0) * 0.08
+        hump = torch.sin(torch.linspace(0.0, torch.pi, frames)).unsqueeze(-1)
+        mirrored = sampled + hump * perpendicular * amplitude
+    mirrored[-1] = sampled[-1]
+    return torch.stack((sampled, mirrored))