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" 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, } 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 future_observation_proxy( image: torch.Tensor, history_xy: torch.Tensor, future_xy: torch.Tensor, ) -> torch.Tensor: """Create a tiny future-camera proxy instead of encoding coordinates as a heatmap. 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. """ 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( 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))