seeker
206 lines
8.1 KiB
Python
206 lines
8.1 KiB
Python
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))
|