569 lines
26 KiB
Python
569 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
r"""Preprocessing for the π0.5 VLA model.
|
|
|
|
Converts a raw robot observation (multi-camera images + language instruction +
|
|
proprioceptive state) into the tensors ``Pi05ForActionPrediction.sample_actions``
|
|
consumes.
|
|
|
|
**The defining π0.5 difference:** the state is not projected by a ``state_proj``
|
|
layer. It is normalized to ``[-1, 1]``, discretized into ``state_num_bins`` bins,
|
|
and serialized into the language prompt::
|
|
|
|
"Task: <instruction>, State: <b0> <b1> ... <bN>;\nAction: "
|
|
|
|
so ``sample_actions`` receives no state tensor at all.
|
|
|
|
**Normalization must precede discretization.** The discretizer bins over
|
|
``[-1, 1]`` and assumes the state is already in that range. Reversed, every bin
|
|
index is wrong and nothing raises.
|
|
|
|
Reference:
|
|
- OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad)
|
|
- OpenPI: openpi/src/openpi/models/pi0_config.py (PaliGemmaTokenizer.tokenize)
|
|
- LeRobot: lerobot/src/lerobot/policies/pi05/processor_pi05.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.models.pi05.config import resolve_excluded_action_indices
|
|
from vllm_omni.diffusion.models.pi05.modeling_pi05 import (
|
|
DEFAULT_IMAGE_RESOLUTION,
|
|
DEFAULT_MAX_TOKEN_LEN,
|
|
DEFAULT_STATE_NUM_BINS,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# LeRobot NormalizerProcessorStep.eps.
|
|
NORM_EPS = 1e-8
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# Image preprocessing (identical to π0 — SigLIP is unchanged in π0.5)
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
def resize_with_pad(
|
|
images: torch.Tensor,
|
|
target_height: int,
|
|
target_width: int,
|
|
mode: str = "bilinear",
|
|
) -> torch.Tensor:
|
|
"""Resize ``(B, C, H, W)`` images to the target shape, preserving aspect
|
|
ratio with -1 padding on the short side.
|
|
|
|
Matches openpi ``image_tools.resize_with_pad_torch`` — the clamp to
|
|
[-1, 1] is what lets the padded region blend with SigLIP-normalized
|
|
pixels without adding signal at the boundary.
|
|
"""
|
|
if images.ndim != 4:
|
|
raise ValueError(f"Expected 4-D (B,C,H,W), got {images.ndim}-D")
|
|
_, _, cur_h, cur_w = images.shape
|
|
ratio = max(cur_w / target_width, cur_h / target_height)
|
|
rh, rw = int(cur_h / ratio), int(cur_w / ratio)
|
|
align_corners = False if mode == "bilinear" else None
|
|
resized = F.interpolate(images, size=(rh, rw), mode=mode, align_corners=align_corners)
|
|
resized = resized.clamp(-1.0, 1.0)
|
|
ph, rem_h = divmod(target_height - rh, 2)
|
|
pw, rem_w = divmod(target_width - rw, 2)
|
|
return F.pad(resized, (pw, pw + rem_w, ph, ph + rem_h), value=-1.0)
|
|
|
|
|
|
def pil_image_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
"""PIL → ``(1, C, H, W)`` float32 in ``[-1, 1]`` (SigLIP normalization)."""
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
arr = np.array(image, dtype=np.float32) / 255.0 * 2.0 - 1.0
|
|
return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
|
|
|
|
|
class Pi05ImageProcessor:
|
|
"""Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor."""
|
|
|
|
def __init__(self, image_size: int = DEFAULT_IMAGE_RESOLUTION[0]):
|
|
self.image_size = image_size
|
|
|
|
def preprocess_single(self, image: Any) -> torch.Tensor:
|
|
"""Convert one LeRobot-domain image to a normalized model tensor.
|
|
|
|
PIL and uint8 inputs have domain ``[0, 255]``. Floating ndarray/tensor
|
|
inputs must be finite and in ``[0, 1]``. Values are converted
|
|
deterministically to ``[-1, 1]``; the pixel content is never used to
|
|
guess its input domain.
|
|
"""
|
|
t = self._to_tensor(image)
|
|
if t.shape[2] != self.image_size or t.shape[3] != self.image_size:
|
|
t = resize_with_pad(t, self.image_size, self.image_size)
|
|
return t
|
|
|
|
def _to_tensor(self, image: Any) -> torch.Tensor:
|
|
if isinstance(image, Image.Image):
|
|
return pil_image_to_tensor(image)
|
|
if isinstance(image, np.ndarray):
|
|
if image.ndim != 3 or image.shape[-1] != 3:
|
|
raise ValueError(f"Expected an HWC ndarray with 3 channels, got shape {image.shape}.")
|
|
t = torch.from_numpy(np.ascontiguousarray(image)).permute(2, 0, 1).unsqueeze(0)
|
|
elif isinstance(image, torch.Tensor):
|
|
t = image.unsqueeze(0) if image.ndim == 3 else image
|
|
if t.ndim != 4 or t.shape[0] != 1 or t.shape[1] != 3:
|
|
raise ValueError(f"Expected a CHW tensor (optionally batched once), got shape {tuple(t.shape)}.")
|
|
else:
|
|
raise TypeError(f"Unsupported image type for π0.5 preprocessing: {type(image)}")
|
|
|
|
if t.dtype == torch.uint8:
|
|
return t.to(dtype=torch.float32) / 255.0 * 2.0 - 1.0
|
|
if not t.is_floating_point():
|
|
raise TypeError(f"π0.5 images must use uint8 or floating dtype, got {t.dtype}.")
|
|
if not torch.isfinite(t).all():
|
|
raise ValueError("Floating π0.5 image contains NaN or Inf.")
|
|
low, high = float(t.amin()), float(t.amax())
|
|
if low < 0.0 or high > 1.0:
|
|
raise ValueError(f"Floating π0.5 images must be in [0, 1]; observed min={low}, max={high}.")
|
|
return t.to(dtype=torch.float32) * 2.0 - 1.0
|
|
|
|
def make_empty_image(self) -> torch.Tensor:
|
|
"""Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot."""
|
|
return torch.full((1, 3, self.image_size, self.image_size), -1.0)
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# Normalization (LeRobot NormalizerProcessorStep)
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
@dataclass(frozen=True)
|
|
class NormStats:
|
|
"""The affine map a checkpoint's statistics define.
|
|
|
|
``mean_std`` carries mean and std; the two range modes carry the lower and
|
|
upper bound (``min``/``max``, or the ``q01``/``q99`` quantiles π0.5 ships).
|
|
"""
|
|
|
|
mode: str
|
|
lower: torch.Tensor
|
|
upper: torch.Tensor
|
|
|
|
|
|
def build_norm_stats(norm_stats: dict | None, key: str) -> NormStats | None:
|
|
"""Read ``norm_stats[key]`` into tensors, or ``None`` when it carries none.
|
|
|
|
The mode comes from the entry. A LeRobot state_dict ships mean, std, min,
|
|
max, q01 and q99 at once, so the statistic names present cannot select it.
|
|
Missing modes default to quantile, which is what LeRobot defaults π0.5's
|
|
STATE and ACTION to (π0 defaults to mean_std).
|
|
"""
|
|
entry = (norm_stats or {}).get(key)
|
|
if not entry:
|
|
return None
|
|
|
|
mode = str(entry.get("mode", "quantile")).lower()
|
|
if mode == "mean_std":
|
|
names = ("mean", "std")
|
|
elif mode == "min_max":
|
|
names = ("min", "max")
|
|
elif mode == "quantile":
|
|
mode, names = "min_max", ("q01", "q99")
|
|
else:
|
|
raise ValueError(f"Unsupported normalization mode for {key!r}: {mode!r}.")
|
|
|
|
bounds = [entry.get(name) for name in names]
|
|
if any(bound is None for bound in bounds):
|
|
raise ValueError(f"Normalization mode {mode!r} for {key!r} requires {names[0]!r} and {names[1]!r}.")
|
|
lower, upper = (torch.as_tensor(bound, dtype=torch.float32) for bound in bounds)
|
|
return NormStats(mode=mode, lower=lower, upper=upper)
|
|
|
|
|
|
def apply_norm(x: torch.Tensor, stats: NormStats | None, *, inverse: bool = False) -> torch.Tensor:
|
|
"""Map raw units to the model's normalized space, or back with ``inverse``.
|
|
|
|
Actions are padded to ``max_action_dim`` while the statistics cover only the
|
|
real width, so the tail passes through untouched. The eps rules are
|
|
LeRobot's: ``mean_std`` always divides by ``std + eps``, and the range modes
|
|
substitute ``eps`` only for an exactly zero range.
|
|
"""
|
|
if stats is None:
|
|
return x
|
|
lower = stats.lower.to(device=x.device, dtype=x.dtype)
|
|
upper = stats.upper.to(device=x.device, dtype=x.dtype)
|
|
valid = lower.shape[0]
|
|
head = x[..., :valid]
|
|
|
|
if stats.mode == "mean_std":
|
|
head = head * upper + lower if inverse else (head - lower) / (upper + NORM_EPS)
|
|
else:
|
|
span = upper - lower
|
|
span = torch.where(span == 0, NORM_EPS, span)
|
|
head = (head + 1.0) * 0.5 * span + lower if inverse else 2.0 * (head - lower) / span - 1.0
|
|
|
|
if valid == x.shape[-1]:
|
|
return head
|
|
return torch.cat([head, x[..., valid:]], dim=-1)
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# State: normalize → discretize → prompt (π0.5's defining path)
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
def _pad_or_truncate(raw_state: Any, width: int) -> np.ndarray:
|
|
"""Zero-pad / truncate a vector to ``(width,)`` float32."""
|
|
if raw_state is None:
|
|
return np.zeros((width,), dtype=np.float32)
|
|
if isinstance(raw_state, torch.Tensor):
|
|
raw_state = raw_state.detach().cpu().numpy()
|
|
state = np.asarray(raw_state, dtype=np.float32).reshape(-1)
|
|
if state.shape[0] < width:
|
|
state = np.pad(state, (0, width - state.shape[0]))
|
|
elif state.shape[0] > width:
|
|
state = state[:width]
|
|
return state.astype(np.float32)
|
|
|
|
|
|
def as_state_vector(raw_state: Any, state_dim: int) -> np.ndarray:
|
|
"""Coerce a request's raw state to ``(state_dim,)`` float32, or raise.
|
|
|
|
LeRobot's ``Pi05PrepareStateTokenizerProcessorStep`` discretizes whatever
|
|
width it is handed and never pads to ``max_state_dim``, so zero-filling a
|
|
7-dim state up to 32 would append 25 state tokens the checkpoint never saw
|
|
in training.
|
|
"""
|
|
if raw_state is None:
|
|
raise ValueError("π0.5 requires a state; there is no default to fall back on.")
|
|
if isinstance(raw_state, torch.Tensor):
|
|
raw_state = raw_state.detach().cpu().numpy()
|
|
state = np.asarray(raw_state, dtype=np.float32)
|
|
if state.ndim == 2 and state.shape[0] == 1:
|
|
state = state[0]
|
|
if state.ndim != 1:
|
|
raise ValueError(f"π0.5 state must be shaped (D,) or (1, D); got {tuple(state.shape)}.")
|
|
if not np.isfinite(state).all():
|
|
raise ValueError("π0.5 state contains NaN or Inf.")
|
|
if state.shape[0] != state_dim:
|
|
raise ValueError(
|
|
f"π0.5 state has {state.shape[0]} dimension(s), but the checkpoint declares "
|
|
f"{state_dim}. The state is serialized into the prompt at its real width, so "
|
|
"padding or truncating it would change the tokens the model sees."
|
|
)
|
|
return state
|
|
|
|
|
|
def discretize_state(state: np.ndarray, *, num_bins: int = DEFAULT_STATE_NUM_BINS) -> np.ndarray:
|
|
"""Discretize a ``[-1, 1]`` state into ``num_bins`` integer bins.
|
|
|
|
Byte-for-byte LeRobot's ``processor_pi05.py``::
|
|
|
|
np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
|
|
|
|
A state below ``-1`` lands in bin ``-1``, and that negative bin is part of
|
|
the contract: the checkpoint was trained with ``" -1"`` in the state prompt
|
|
for those dimensions. Clipping it to ``0`` changes the tokens the model
|
|
sees, so this must not clip.
|
|
|
|
``bins`` stays float64, matching LeRobot's default ``linspace`` dtype, so
|
|
boundary values fall on the same side.
|
|
|
|
**Assumes the state is already normalized** — see the module docstring.
|
|
"""
|
|
bins = np.linspace(-1.0, 1.0, num_bins + 1)[:-1]
|
|
return (np.digitize(np.asarray(state, dtype=np.float32), bins=bins) - 1).astype(np.int64)
|
|
|
|
|
|
def build_pi05_prompt(
|
|
*,
|
|
task: str,
|
|
normalized_state: np.ndarray,
|
|
state_num_bins: int = DEFAULT_STATE_NUM_BINS,
|
|
) -> str:
|
|
"""Build the π0.5 prompt: instruction + serialized discretized state.
|
|
|
|
Matches LeRobot's ``Pi05PrepareStateTokenizerProcessorStep``, including the
|
|
task cleanup (``strip``, ``_`` → space, newline → space), the exact template
|
|
and the state values it serializes. The template already ends in a newline,
|
|
so — unlike π0 — there is no separate newline-appending step.
|
|
|
|
``normalized_state`` comes from ``Pi05ForActionPrediction._normalize_state``.
|
|
"""
|
|
cleaned_task = (task or "").strip().replace("_", " ").replace("\n", " ")
|
|
bins = discretize_state(normalized_state, num_bins=state_num_bins)
|
|
state_str = " ".join(str(int(x)) for x in bins.tolist())
|
|
return f"Task: {cleaned_task}, State: {state_str};\nAction: "
|
|
|
|
|
|
def tokenize_prompt(tokenizer, text: str, max_token_len: int = DEFAULT_MAX_TOKEN_LEN):
|
|
"""Return ``(input_ids, attention_mask)`` lists, length exactly ``max_token_len``.
|
|
|
|
``padding="max_length"`` is what makes the prefix a constant shape: the text
|
|
segment is always ``max_token_len`` tokens regardless of the instruction, so
|
|
only ``attention_mask.sum()`` varies per request.
|
|
"""
|
|
enc = tokenizer(
|
|
text,
|
|
padding="max_length",
|
|
max_length=max_token_len,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors=None,
|
|
)
|
|
return list(enc["input_ids"]), list(enc["attention_mask"])
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# Relative actions
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
class Pi05RelativeActions:
|
|
"""The relative/absolute action transform, as a single paired object.
|
|
|
|
LeRobot builds one ``RelativeActionsProcessorStep`` and hands *the same
|
|
instance* to ``AbsoluteActionsProcessorStep``; the two directions must
|
|
agree on ``enabled`` and on which dimensions are excluded, so they are one
|
|
object here too.
|
|
|
|
**Deviation from LeRobot, on purpose.** LeRobot's step keeps the reference
|
|
state on ``self`` between the pre- and post-pass. That is safe for a
|
|
single-threaded training loop and unsafe for a server: two in-flight
|
|
requests would share one reference state and silently corrupt each other's
|
|
actions. Here the state is passed explicitly to :meth:`to_absolute`, so the
|
|
object stays immutable after construction and is safe to share across
|
|
requests.
|
|
|
|
Transform (LeRobot / OpenPI): ``relative = action - state`` on the way in,
|
|
``absolute = relative + state`` on the way out, applied only to the
|
|
dimensions *not* named in ``exclude_joints``. Gripper open/close is an
|
|
absolute command, which is why it is excluded by default.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
enabled: bool,
|
|
exclude_joints: list[str] | None = None,
|
|
action_names: list[str] | None = None,
|
|
max_action_dim: int = 32,
|
|
):
|
|
self.enabled = bool(enabled)
|
|
self.exclude_joints = list(exclude_joints or [])
|
|
self.action_names = list(action_names) if action_names else None
|
|
self.max_action_dim = int(max_action_dim)
|
|
|
|
# Boolean mask over action dims: True = this dim is relative to state.
|
|
mask = np.ones((self.max_action_dim,), dtype=bool)
|
|
if self.enabled:
|
|
for idx in resolve_excluded_action_indices(self.exclude_joints, self.action_names):
|
|
if 0 <= idx < self.max_action_dim:
|
|
mask[idx] = False
|
|
# Padding beyond the real action dimensions carries no signal;
|
|
# leaving it "relative" would add state noise into dead channels.
|
|
if self.action_names:
|
|
mask[len(self.action_names) :] = False
|
|
else:
|
|
mask[:] = False
|
|
self.relative_mask = mask
|
|
|
|
@property
|
|
def num_relative_dims(self) -> int:
|
|
return int(self.relative_mask.sum())
|
|
|
|
def _state_row(self, state: Any, device, dtype) -> torch.Tensor:
|
|
"""Raw state → ``(B, max_action_dim)`` aligned with the action dims.
|
|
|
|
Accepts one state for the whole batch (``(D,)``, the serving path,
|
|
where the pipeline runs B=1) or one state per sample (``(B, D)``).
|
|
LeRobot caches the batched state and shifts each sample by its own row,
|
|
so the per-sample form has to be honoured: ``_pad_or_truncate``
|
|
flattens, which would silently reduce ``(B, D)`` to sample 0's state and
|
|
apply it to every sample — wrong answers, no error.
|
|
"""
|
|
if isinstance(state, torch.Tensor):
|
|
state = state.detach().cpu().numpy()
|
|
arr = np.asarray(0.0 if state is None else state, dtype=np.float32)
|
|
if arr.ndim > 2:
|
|
raise ValueError(f"Expected state shaped (D,) or (B, D), got {tuple(arr.shape)}")
|
|
rows = arr if arr.ndim == 2 else arr[None, :]
|
|
padded = np.stack([_pad_or_truncate(row, self.max_action_dim) for row in rows])
|
|
return torch.as_tensor(padded, device=device, dtype=dtype)
|
|
|
|
def to_relative(self, actions: torch.Tensor, state: Any) -> torch.Tensor:
|
|
"""``absolute → relative``. Input side (step 3).
|
|
|
|
Not used on the inference path — there are no input actions to convert
|
|
at serving time — but it is what the transform *means*, and the parity
|
|
test exercises it as the inverse of :meth:`to_absolute`.
|
|
"""
|
|
if not self.enabled:
|
|
return actions
|
|
return self._shift(actions, state, sign=-1.0)
|
|
|
|
def to_absolute(self, actions: torch.Tensor, state: Any) -> torch.Tensor:
|
|
"""``relative → absolute``. Output side (step 2 of the post-pipeline).
|
|
|
|
``state`` must be the **raw** state — the same one the model was given
|
|
before normalization — because relative actions live in raw action space.
|
|
"""
|
|
if not self.enabled:
|
|
return actions
|
|
return self._shift(actions, state, sign=+1.0)
|
|
|
|
def _shift(self, actions: torch.Tensor, state: Any, *, sign: float) -> torch.Tensor:
|
|
if actions.ndim != 3:
|
|
raise ValueError(f"Expected actions shaped (B, horizon, action_dim), got {tuple(actions.shape)}")
|
|
action_dim = actions.shape[-1]
|
|
if action_dim != self.max_action_dim:
|
|
raise ValueError(
|
|
f"Action dim {action_dim} does not match max_action_dim={self.max_action_dim}; "
|
|
"the relative mask would be misaligned."
|
|
)
|
|
state_row = self._state_row(state, actions.device, actions.dtype) # (1, D)
|
|
mask = torch.as_tensor(self.relative_mask, device=actions.device)
|
|
delta = torch.where(mask, state_row, torch.zeros_like(state_row))
|
|
# Broadcast over the action horizon: every step of the chunk is
|
|
# expressed relative to the same current state.
|
|
return actions + sign * delta[:, None, :]
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# Model input assembly
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
def _extract_images(robot_obs: dict, config) -> dict[str, Any]:
|
|
"""Pull a ``{feature_key: image}`` map out of a raw robot obs.
|
|
|
|
This is functional step 1 (``rename_observations``): keys are translated
|
|
through ``config.image_key_map`` so serving wire names map onto the
|
|
checkpoint's ``input_features`` identities.
|
|
"""
|
|
images = robot_obs.get("images")
|
|
if not isinstance(images, dict):
|
|
images = {k: v for k, v in robot_obs.items() if _is_image_like(v)}
|
|
key_map = config.image_key_map or {}
|
|
return {key_map.get(k, k): v for k, v in images.items() if _is_image_like(v)}
|
|
|
|
|
|
def _is_image_like(value: Any) -> bool:
|
|
if isinstance(value, (Image.Image, torch.Tensor)):
|
|
return True
|
|
if isinstance(value, np.ndarray):
|
|
return value.ndim >= 3
|
|
if isinstance(value, (list, tuple)):
|
|
try:
|
|
return np.asarray(value).ndim >= 3
|
|
except Exception: # noqa: BLE001
|
|
return False
|
|
return False
|
|
|
|
|
|
def _assemble_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device, state_norm: NormStats | None):
|
|
"""Convert a raw robot observation into ``sample_actions`` inputs.
|
|
|
|
Returns ``(images, image_masks, lang_tokens, lang_masks)`` — note there is
|
|
**no state tensor**: π0.5 carries the state inside ``lang_tokens``.
|
|
|
|
Camera slots follow ``config.image_feature_keys`` order. A missing key keeps
|
|
its semantic slot and receives an empty image with a false mask. The output
|
|
always contains exactly ``max_cameras`` slots for the deployed model.
|
|
"""
|
|
image_size = int(config.image_resolution[0])
|
|
img_proc = Pi05ImageProcessor(image_size=image_size)
|
|
max_cameras = config.max_cameras
|
|
|
|
feature_keys = config.image_feature_keys or []
|
|
obs_images = _extract_images(robot_obs, config)
|
|
if not feature_keys:
|
|
# No declared camera order — fall back to whatever the obs provides,
|
|
# preserving insertion order.
|
|
feature_keys = list(obs_images.keys())[:max_cameras]
|
|
|
|
camera_keys = feature_keys[:max_cameras]
|
|
if all(obs_images.get(key) is None for key in camera_keys):
|
|
raise ValueError("π0.5 observation must provide at least one configured camera image.")
|
|
if not config.image_feature_keys and len(obs_images) > max_cameras:
|
|
raise ValueError(
|
|
f"π0.5 observation provides {len(obs_images)} cameras, which exceeds max_cameras={max_cameras}."
|
|
)
|
|
|
|
images: list[torch.Tensor] = []
|
|
image_masks: list[torch.Tensor] = []
|
|
for key in camera_keys:
|
|
image = obs_images.get(key)
|
|
if image is None:
|
|
tensor = img_proc.make_empty_image().to(device=device)
|
|
mask = False
|
|
else:
|
|
tensor = img_proc.preprocess_single(image).to(device=device)
|
|
mask = True
|
|
images.append(tensor)
|
|
image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device))
|
|
|
|
# A checkpoint may declare fewer named slots than the serving graph.
|
|
while len(images) < max_cameras:
|
|
images.append(img_proc.make_empty_image().to(device=device))
|
|
image_masks.append(torch.tensor([False], dtype=torch.bool, device=device))
|
|
|
|
# Normalize and serialize the state into the tokenized prompt.
|
|
raw_state = as_state_vector(robot_obs.get("state"), config.state_dim)
|
|
normalized_state = apply_norm(torch.from_numpy(raw_state), state_norm).numpy()
|
|
prompt = build_pi05_prompt(
|
|
task=robot_obs.get("prompt", "") or "",
|
|
normalized_state=normalized_state,
|
|
state_num_bins=config.state_num_bins,
|
|
)
|
|
ids, attn = tokenize_prompt(tokenizer, prompt, config.tokenizer_max_length)
|
|
lang_tokens = torch.tensor([ids], dtype=torch.long, device=device)
|
|
lang_masks = torch.tensor([attn], dtype=torch.bool, device=device)
|
|
|
|
return images, image_masks, lang_tokens, lang_masks
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
# The interface the pipeline uses
|
|
# ──────────────────────────────────────────────────────────────────────
|
|
class Pi05Processor:
|
|
"""Everything between the OpenPI wire and the model.
|
|
|
|
Owns the normalization statistics and the relative-action transform, so the
|
|
pipeline holds no preprocessing state of its own.
|
|
"""
|
|
|
|
def __init__(self, config, tokenizer, device: torch.device):
|
|
self.config = config
|
|
self.tokenizer = tokenizer
|
|
self.device = device
|
|
self._state_norm = build_norm_stats(config.norm_stats, "state")
|
|
self._action_norm = build_norm_stats(config.norm_stats, "action")
|
|
self.relative_actions = Pi05RelativeActions(
|
|
enabled=config.use_relative_actions,
|
|
exclude_joints=config.relative_exclude_joints,
|
|
action_names=config.action_feature_names,
|
|
max_action_dim=config.max_action_dim,
|
|
)
|
|
if self._action_norm is None:
|
|
logger.info("π0.5: no action normalization stats; returned actions stay in normalized space.")
|
|
if self.relative_actions.enabled:
|
|
logger.info(
|
|
"π0.5: relative actions enabled — %d of %d action dims are state-relative (excluded joints: %s).",
|
|
self.relative_actions.num_relative_dims,
|
|
config.max_action_dim,
|
|
config.relative_exclude_joints,
|
|
)
|
|
|
|
def build_model_inputs(self, robot_obs: dict):
|
|
"""Raw observation → ``sample_actions`` inputs.
|
|
|
|
No state tensor comes back: π0.5 carries the state inside ``lang_tokens``.
|
|
"""
|
|
return _assemble_model_inputs(robot_obs, self.config, self.tokenizer, self.device, self._state_norm)
|
|
|
|
def build_model_outputs(self, actions: torch.Tensor, robot_obs: dict) -> np.ndarray:
|
|
"""Model output → the action chunk the robot receives.
|
|
|
|
``AbsoluteActionsProcessorStep`` runs after unnormalization, because a
|
|
relative-action checkpoint's statistics are computed in relative space.
|
|
The raw state is the one the prompt encoded, before normalization.
|
|
"""
|
|
actions = apply_norm(actions, self._action_norm, inverse=True)
|
|
if self.relative_actions.enabled:
|
|
actions = self.relative_actions.to_absolute(actions, robot_obs.get("state"))
|
|
return actions.squeeze(0)[..., : self.config.action_dim].float().cpu().numpy()
|