diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index 45f1d81b..e46b8676 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -108,16 +108,22 @@ def __repr__(self) -> str: component_registry.register("smolvla_resize", "physicalai.inference.preprocessors.ResizeSmolVLA") component_registry.register("new_line", "physicalai.inference.preprocessors.NewLinePreprocessor") component_registry.register("hf_tokenizer", "physicalai.inference.preprocessors.HFTokenizer") +component_registry.register("molmoact2", "physicalai.inference.preprocessors.MolmoAct2Preprocessor") +component_registry.register("molmoact2_inputs", "physicalai.inference.preprocessors.MolmoAct2ModelInputs") component_registry.register("ov_tokenizer", "physicalai.inference.preprocessors.OVTokenizer") component_registry.register("pi05", "physicalai.inference.preprocessors.Pi05Preprocessor") component_registry.register("rldx1", "physicalai.inference.preprocessors.Rldx1Preprocessor") component_registry.register("rldx1_token_composer", "physicalai.inference.preprocessors.Rldx1TokenComposer") component_registry.register("rldx1_rope", "physicalai.inference.preprocessors.Rldx1RopePreprocessor") component_registry.register("to_float_tensor", "physicalai.inference.preprocessors.ToFloatTensorPreprocessor") +component_registry.register("molmoact2_pre", "physicalai.inference.preprocessors.MolmoAct2Preprocessor") +component_registry.register("joint_frame_preprocess", "physicalai.inference.preprocessors.JointFramePreprocessor") # Postprocessors component_registry.register("denormalize", "physicalai.inference.postprocessors.StatsDenormalizer") component_registry.register("action_chunk_trimmer", "physicalai.inference.postprocessors.ActionChunkTrimmer") +component_registry.register("molmoact2_postprocess", "physicalai.inference.postprocessors.MolmoAct2Postprocessor") +component_registry.register("joint_frame_postprocess", "physicalai.inference.postprocessors.JointFramePostprocessor") # Callbacks component_registry.register("latency_monitor", "physicalai.inference.callbacks.LatencyMonitor") diff --git a/src/physicalai/inference/postprocessors/__init__.py b/src/physicalai/inference/postprocessors/__init__.py index 64a45d42..f38feab6 100644 --- a/src/physicalai/inference/postprocessors/__init__.py +++ b/src/physicalai/inference/postprocessors/__init__.py @@ -9,11 +9,15 @@ from physicalai.inference.postprocessors.action_chunk_trimmer import ActionChunkTrimmer from physicalai.inference.postprocessors.action_normalizer import ActionNormalizer from physicalai.inference.postprocessors.base import Postprocessor +from physicalai.inference.postprocessors.joint_frame import JointFramePostprocessor +from physicalai.inference.postprocessors.molmoact2 import MolmoAct2Postprocessor from physicalai.inference.postprocessors.stats_denormalizer import StatsDenormalizer __all__ = [ "ActionChunkTrimmer", "ActionNormalizer", + "JointFramePostprocessor", + "MolmoAct2Postprocessor", "Postprocessor", "StatsDenormalizer", ] diff --git a/src/physicalai/inference/postprocessors/joint_frame.py b/src/physicalai/inference/postprocessors/joint_frame.py new file mode 100644 index 00000000..55677410 --- /dev/null +++ b/src/physicalai/inference/postprocessors/joint_frame.py @@ -0,0 +1,91 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Joint-frame action postprocessing.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from typing_extensions import override + +from physicalai.inference.postprocessors.base import Postprocessor + +if TYPE_CHECKING: + from collections.abc import Sequence + + +class JointFrameTransform: + """Apply an invertible affine transform to leading joint values.""" + + def __init__(self, *, signs: Sequence[float], offsets: Sequence[float]) -> None: + """Store the joint signs and offsets. + + Raises: + ValueError: If signs and offsets differ in length or a sign is not +/-1. + """ + if len(signs) != len(offsets): + msg = f"signs ({len(signs)}) and offsets ({len(offsets)}) must match" + raise ValueError(msg) + if any(sign not in {-1.0, 1.0} for sign in signs): + msg = "Joint frame transform signs must be either -1 or 1." + raise ValueError(msg) + self._signs = np.asarray(signs, dtype=np.float32) + self._offsets = np.asarray(offsets, dtype=np.float32) + + def forward(self, values: np.ndarray) -> np.ndarray: + """Apply ``sign * value + offset`` to leading joint values. + + Returns: + A transformed copy of ``values``. + """ + return self._apply(values, inverse=False) + + def inverse(self, values: np.ndarray) -> np.ndarray: + """Apply ``sign * (value - offset)`` to leading joint values. + + Returns: + An inverse-transformed copy of ``values``. + """ + return self._apply(values, inverse=True) + + def _apply(self, values: np.ndarray, *, inverse: bool) -> np.ndarray: + count = min(self._signs.size, values.shape[-1]) + output = np.array(values, copy=True) + joints = values[..., :count] + output[..., :count] = ( + self._signs[:count] * (joints - self._offsets[:count]) + if inverse + else self._signs[:count] * joints + self._offsets[:count] + ) + return output + + +class JointFramePostprocessor(Postprocessor): + """Map one output feature from checkpoint to robot joint coordinates.""" + + def __init__(self, *, feature: str, signs: Sequence[float], offsets: Sequence[float]) -> None: + """Configure the feature and calibration frame.""" + self._feature = feature + self._transform = JointFrameTransform(signs=signs, offsets=offsets) + + @override + def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + """Transform the configured feature while preserving all other outputs. + + Returns: + A shallow copy with the transformed feature. + + Raises: + ValueError: If the configured feature is absent. + """ + if self._feature not in outputs: + msg = f"Joint frame postprocessor expected feature {self._feature!r}" + raise ValueError(msg) + result = dict(outputs) + result[self._feature] = self._transform.inverse(np.asarray(outputs[self._feature])) + return result + + +__all__ = ["JointFramePostprocessor"] diff --git a/src/physicalai/inference/postprocessors/molmoact2/__init__.py b/src/physicalai/inference/postprocessors/molmoact2/__init__.py new file mode 100644 index 00000000..f686e05c --- /dev/null +++ b/src/physicalai/inference/postprocessors/molmoact2/__init__.py @@ -0,0 +1,8 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""MolmoAct2 inference postprocessors.""" + +from physicalai.inference.postprocessors.molmoact2.processor import MolmoAct2Postprocessor + +__all__ = ["MolmoAct2Postprocessor"] diff --git a/src/physicalai/inference/postprocessors/molmoact2/processor.py b/src/physicalai/inference/postprocessors/molmoact2/processor.py new file mode 100644 index 00000000..67b6ce01 --- /dev/null +++ b/src/physicalai/inference/postprocessors/molmoact2/processor.py @@ -0,0 +1,69 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy postprocessing for MolmoAct2 inference.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +from typing_extensions import override + +from physicalai.inference.constants import ACTION +from physicalai.inference.postprocessors.base import Postprocessor +from physicalai.inference.postprocessors.stats_denormalizer import StatsDenormalizer +from physicalai.inference.preprocessors.molmoact2 import normalization_stats + + +class MolmoAct2Postprocessor(Postprocessor): + """Clamp and denormalize MolmoAct2 actions.""" + + def __init__( + self, + *, + action_key: str, + action_stats: dict[str, Any] | None = None, + normalization_mode: str = "QUANTILES", + ) -> None: + """Store action postprocessing settings. + + Args: + action_key: Adapter output key containing the action tensor. + action_stats: Statistics used to denormalize actions. + normalization_mode: Normalization strategy used during training. + """ + self._action_key = action_key + self.denormalizer = ( + StatsDenormalizer( + stats={ACTION: normalization_stats(action_stats)}, + mode=normalization_mode.lower(), + features=[ACTION], + ) + if action_stats + else None + ) + + @override + def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + """Postprocess the model action output. + + Returns: + Outputs with the canonical denormalized action. + + Raises: + ValueError: If the configured action output is absent. + """ + result = dict(outputs) + if self._action_key not in result: + msg = f"MolmoAct2 postprocessor expected action key {self._action_key!r}" + raise ValueError(msg) + action = result.pop(self._action_key) + action = np.clip(np.asarray(action), -1.0, 1.0) + if self.denormalizer is not None: + action = self.denormalizer({ACTION: action})[ACTION] + result[ACTION] = action + return result + + +__all__ = ["MolmoAct2Postprocessor"] diff --git a/src/physicalai/inference/postprocessors/stats_denormalizer.py b/src/physicalai/inference/postprocessors/stats_denormalizer.py index 53fd6428..7bc2888a 100644 --- a/src/physicalai/inference/postprocessors/stats_denormalizer.py +++ b/src/physicalai/inference/postprocessors/stats_denormalizer.py @@ -186,21 +186,23 @@ def _denormalize( Returns: Denormalized array. """ + transformed = tensor if mode == "mean_std": mean = stats["mean"] std = stats["std"] - return tensor * std + mean + transformed = tensor * std + mean - if mode == "min_max": + elif mode == "min_max": min_val = stats["min"] max_val = stats["max"] - return (tensor + 1.0) / 2.0 * (max_val - min_val) + min_val + transformed = (tensor + 1.0) / 2.0 * (max_val - min_val) + min_val - if mode == "quantiles": + elif mode == "quantiles": q01 = stats["q01"] q99 = stats["q99"] denom = q99 - q01 denom = np.where(denom == 0, _EPS, denom) - return (tensor + 1.0) * denom / 2.0 + q01 + transformed = (tensor + 1.0) * denom / 2.0 + q01 - return tensor + mask = stats.get("mask") + return np.where(mask.astype(np.bool_), transformed, tensor) if mask is not None else transformed diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index c006b908..0485de9d 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -9,7 +9,9 @@ from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer +from physicalai.inference.preprocessors.joint_frame import JointFramePreprocessor from physicalai.inference.preprocessors.lambda_processor import LambdaPreprocessor +from physicalai.inference.preprocessors.molmoact2 import MolmoAct2ModelInputs, MolmoAct2Preprocessor from physicalai.inference.preprocessors.new_line import NewLinePreprocessor from physicalai.inference.preprocessors.ov_tokenizer import OVTokenizer from physicalai.inference.preprocessors.pi05 import Pi05Preprocessor @@ -23,7 +25,10 @@ __all__ = [ "HFTokenizer", + "JointFramePreprocessor", "LambdaPreprocessor", + "MolmoAct2ModelInputs", + "MolmoAct2Preprocessor", "NewLinePreprocessor", "OVTokenizer", "Pi05Preprocessor", diff --git a/src/physicalai/inference/preprocessors/joint_frame.py b/src/physicalai/inference/preprocessors/joint_frame.py new file mode 100644 index 00000000..08ae19a6 --- /dev/null +++ b/src/physicalai/inference/preprocessors/joint_frame.py @@ -0,0 +1,50 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Joint-frame observation preprocessing.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +from typing_extensions import override + +from physicalai.inference.postprocessors.joint_frame import JointFrameTransform +from physicalai.inference.preprocessors.base import Preprocessor + +if TYPE_CHECKING: + from collections.abc import Sequence + + +class JointFramePreprocessor(Preprocessor): + """Map one observation feature from robot to checkpoint joint coordinates.""" + + def __init__(self, *, feature: str, signs: Sequence[float], offsets: Sequence[float]) -> None: + """Configure the feature and calibration frame.""" + self._feature = feature + self._transform = JointFrameTransform(signs=signs, offsets=offsets) + + @override + def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Transform the configured feature while preserving all other inputs. + + Returns: + A shallow copy with the transformed feature. + """ + key = self._resolve_key(inputs) + outputs = dict(inputs) + outputs[key] = self._transform.forward(np.asarray(inputs[key])) + return outputs + + def _resolve_key(self, inputs: dict[str, Any]) -> str: + if self._feature in inputs: + return self._feature + observation_key = f"observation.{self._feature}" + if observation_key in inputs: + return observation_key + msg = f"Joint frame preprocessor expected feature {self._feature!r}" + raise ValueError(msg) + + +__all__ = ["JointFramePreprocessor"] diff --git a/src/physicalai/inference/preprocessors/molmoact2/__init__.py b/src/physicalai/inference/preprocessors/molmoact2/__init__.py new file mode 100644 index 00000000..42e89285 --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2/__init__.py @@ -0,0 +1,13 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""MolmoAct2 inference preprocessors.""" + +from physicalai.inference.preprocessors.molmoact2.inputs import MolmoAct2ModelInputs +from physicalai.inference.preprocessors.molmoact2.processor import MolmoAct2Preprocessor, normalization_stats + +__all__ = [ + "MolmoAct2ModelInputs", + "MolmoAct2Preprocessor", + "normalization_stats", +] diff --git a/src/physicalai/inference/preprocessors/molmoact2/image.py b/src/physicalai/inference/preprocessors/molmoact2/image.py new file mode 100644 index 00000000..920bb67c --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2/image.py @@ -0,0 +1,102 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy image patchification for MolmoAct2 inference.""" + +from __future__ import annotations + +import numpy as np + +_IMAGE_NDIM = 4 +_NUM_CHANNELS = 3 + + +class MolmoAct2ImageProcessor: + """Normalize and patchify pre-resized BCHW images.""" + + def __init__( + self, + *, + crop_mode: str = "resize", + size: dict[str, int] | None = None, + patch_size: int = 14, + pooling_size: list[int] | tuple[int, int] = (2, 2), + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + ) -> None: + """Store image settings and precompute pooling indices.""" + size = size or {"height": 378, "width": 378} + self.crop_mode = crop_mode + self.height = int(size["height"]) + self.width = int(size["width"]) + self.patch_size = int(patch_size) + self.pool_h, self.pool_w = (int(pooling_size[0]), int(pooling_size[1])) + self.image_mean = image_mean or [0.5, 0.5, 0.5] + self.image_std = image_std or [0.5, 0.5, 0.5] + self._pooling, self.pooled_h, self.pooled_w = self._pooling_indices() + + def __call__(self, images: np.ndarray) -> dict[str, np.ndarray]: + """Return patches, pooling indices, grids, and crop counts. + + Returns: + Model-ready image arrays and layout metadata. + + Raises: + ValueError: If image shape, size, or dtype is unsupported. + NotImplementedError: If crop mode is not ``resize``. + """ + images = np.asarray(images) + if images.ndim != _IMAGE_NDIM or images.shape[1] != _NUM_CHANNELS: + msg = f"Expected images of shape (M, 3, H, W), got {images.shape}." + raise ValueError(msg) + if images.shape[2:] != (self.height, self.width): + msg = f"Expected images of size {(self.height, self.width)}, got {images.shape[2:]}." + raise ValueError(msg) + if self.crop_mode != "resize": + msg = f"MolmoAct2ImageProcessor only supports crop_mode='resize', got {self.crop_mode!r}." + raise NotImplementedError(msg) + if images.dtype not in {np.dtype(np.float16), np.dtype(np.float32)}: + msg = f"Expected images of dtype float16 or float32, got {images.dtype}." + raise ValueError(msg) + + count = images.shape[0] + mean = np.asarray(self.image_mean, dtype=images.dtype).reshape(1, 3, 1, 1) + std = np.asarray(self.image_std, dtype=images.dtype).reshape(1, 3, 1, 1) + pixel_values = self._patchify((images - mean) / std) + pooling = np.tile(self._pooling, (count, 1)) + grid = np.asarray([self.pooled_h, self.pooled_w, 0, 0], dtype=np.int64) + return { + "pixel_values": pixel_values, + "image_token_pooling": pooling, + "image_grids": np.tile(grid, (count, 1)), + "image_num_crops": np.ones(count, dtype=np.int64), + } + + def _patchify(self, pixels: np.ndarray) -> np.ndarray: + count, channels, height, width = pixels.shape + patch = self.patch_size + if height % patch or width % patch: + msg = f"Image size {(height, width)} must be divisible by patch_size={patch}." + raise ValueError(msg) + pixels = pixels.transpose(0, 2, 3, 1) + pixels = pixels.reshape(count, height // patch, patch, width // patch, patch, channels) + return pixels.transpose(0, 1, 3, 2, 4, 5).reshape(count, -1, patch * patch * channels) + + def _pooling_indices(self) -> tuple[np.ndarray, int, int]: + patch_h = self.height // self.patch_size + patch_w = self.width // self.patch_size + pooled_h = (patch_h + self.pool_h - 1) // self.pool_h + pooled_w = (patch_w + self.pool_w - 1) // self.pool_w + pad_h = pooled_h * self.pool_h - patch_h + pad_w = pooled_w * self.pool_w - patch_w + indices = np.arange(patch_h * patch_w, dtype=np.int64).reshape(patch_h, patch_w) + indices = np.pad( + indices, + ((pad_h // 2, (pad_h + 1) // 2), (pad_w // 2, (pad_w + 1) // 2)), + constant_values=-1, + ) + indices = indices.reshape(pooled_h, self.pool_h, pooled_w, self.pool_w) + return indices.transpose(0, 2, 1, 3).reshape(-1, self.pool_h * self.pool_w), pooled_h, pooled_w + + +__all__ = ["MolmoAct2ImageProcessor"] diff --git a/src/physicalai/inference/preprocessors/molmoact2/inputs.py b/src/physicalai/inference/preprocessors/molmoact2/inputs.py new file mode 100644 index 00000000..c2db090e --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2/inputs.py @@ -0,0 +1,471 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy model-input assembly for MolmoAct2 inference. + +Mirrors the PyTorch ``build_model_inputs`` used during training/export so the +exported OpenVINO graph receives identical, fully-prepared tensors. Turns a +tokenized prompt (with ``<|image|>`` placeholders) and patchified images into +``input_ids`` (placeholders expanded), ``attention_mask``, ``token_type_ids``, +per-example batched ``images``, ``token_pooling`` and ``action_dim_is_pad``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from typing_extensions import override + +from physicalai.inference.constants import IMAGES, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.preprocessors.base import Preprocessor + +from .image import MolmoAct2ImageProcessor + +_PACKED_IMAGE_NDIM = 5 + + +@dataclass +class MolmoAct2InputConfig: + """Token ids and layout flags needed to assemble MolmoAct2 model inputs.""" + + pad_token_id: int + image_placeholder_token_id: int + image_patch_id: int + image_start_token_id: int + image_end_token_id: int + image_col_id: int | None = None + low_res_image_start_token_id: int | None = None + frame_start_token_id: int | None = None + frame_end_token_id: int | None = None + image_low_res_id: int | None = None + image_use_col_tokens: bool = True + use_single_crop_col_tokens: bool | None = False + use_single_crop_start_token: bool = True + max_action_dim: int = 32 + env_action_dim: int = 0 + image_token_ids: list[int] | None = None + + def __post_init__(self) -> None: + """Collect the configured image token identifiers.""" + ids = [ + self.image_patch_id, + self.image_col_id, + self.image_start_token_id, + self.low_res_image_start_token_id, + self.frame_start_token_id, + self.image_end_token_id, + self.frame_end_token_id, + self.image_low_res_id, + ] + if self.image_token_ids is None: + self.image_token_ids = [int(token_id) for token_id in ids if token_id is not None] + + +def _image_token_ids_for_grid(config: MolmoAct2InputConfig, grid: np.ndarray) -> list[int]: + """Expand a single image grid into its sequence of image token ids. + + Returns: + Ordered image token identifiers for the grid. + """ + resized_h, resized_w, height, width = (int(x) for x in np.asarray(grid).reshape(-1)[:4].tolist()) + + image_patch_id = int(config.image_patch_id) + image_start_token_id = int(config.image_start_token_id) + image_end_token_id = int(config.image_end_token_id) + image_col_id = None if config.image_col_id is None else int(config.image_col_id) + low_res_start_id = ( + int(config.low_res_image_start_token_id) + if config.low_res_image_start_token_id is not None + else image_start_token_id + ) + + image_use_col_tokens = bool(config.image_use_col_tokens) + use_single_crop_col_tokens = ( + image_use_col_tokens if config.use_single_crop_col_tokens is None else bool(config.use_single_crop_col_tokens) + ) + use_single_crop_start_token = bool(config.use_single_crop_start_token) + + def make_rows(num_rows: int, num_cols: int, *, use_col: bool) -> list[int]: + row = [image_patch_id] * num_cols + if use_col and image_col_id is not None: + row += [image_col_id] + return row * num_rows + + if height == 0 or width == 0: + return [ + image_start_token_id, + *make_rows(resized_h, resized_w, use_col=use_single_crop_col_tokens), + image_end_token_id, + ] + + high_res = [image_start_token_id, *make_rows(height, width, use_col=image_use_col_tokens), image_end_token_id] + low_start = low_res_start_id if use_single_crop_start_token else image_start_token_id + low_res = [low_start, *make_rows(resized_h, resized_w, use_col=use_single_crop_col_tokens), image_end_token_id] + return low_res + high_res + + +def _build_token_type_ids( + config: MolmoAct2InputConfig, input_ids: np.ndarray, attention_mask: np.ndarray +) -> np.ndarray | None: + """Mark image tokens (1) vs. text tokens (0), respecting the attention mask. + + Returns: + Image-token indicators, or ``None`` when no image tokens are configured. + """ + image_token_ids = config.image_token_ids + if not image_token_ids: + return None + token_set = np.asarray(image_token_ids, dtype=input_ids.dtype) + is_image = np.isin(input_ids, token_set).astype(np.int64) + return is_image * attention_mask.astype(np.int64) + + +def expand_image_placeholders( + *, + config: MolmoAct2InputConfig, + input_ids: np.ndarray, + attention_mask: np.ndarray, + image_grids: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: + """Replace each ``<|image|>`` placeholder with its expanded image token ids. + + Args: + config: Token identifiers and image layout configuration. + input_ids: Padded prompt token identifiers. + attention_mask: Mask identifying valid prompt tokens. + image_grids: Image grid dimensions in placeholder order. + + Returns: + Expanded token identifiers, attention mask, and optional token type identifiers. + + Raises: + ValueError: If there are fewer image grids than image placeholders. + """ + if int(image_grids.shape[0]) == 0: + return input_ids, attention_mask, _build_token_type_ids(config, input_ids, attention_mask) + + placeholder_id = int(config.image_placeholder_token_id) + + expanded_rows: list[list[int]] = [] + expanded_widths: list[int] = [] + grid_idx = 0 + for batch_idx in range(int(input_ids.shape[0])): + valid = attention_mask[batch_idx].astype(bool) + expanded: list[int] = [] + for token in input_ids[batch_idx][valid].tolist(): + token_int = int(token) + if token_int == placeholder_id: + if grid_idx >= int(image_grids.shape[0]): + msg = "Not enough image grids to expand all <|image|> placeholders." + raise ValueError(msg) + expanded.extend(_image_token_ids_for_grid(config, image_grids[grid_idx])) + grid_idx += 1 + else: + expanded.append(token_int) + expanded_rows.append(expanded) + expanded_widths.append(len(expanded) + int((~valid).sum())) + + max_len = max(expanded_widths, default=1) + out_ids = np.full((len(expanded_rows), max_len), config.pad_token_id, dtype=input_ids.dtype) + out_mask = np.zeros((len(expanded_rows), max_len), dtype=attention_mask.dtype) + for batch_idx, row in enumerate(expanded_rows): + if not row: + continue + row_arr = np.asarray(row, dtype=input_ids.dtype) + out_ids[batch_idx, : row_arr.size] = row_arr + out_mask[batch_idx, : row_arr.size] = 1 + + return out_ids, out_mask, _build_token_type_ids(config, out_ids, out_mask) + + +@dataclass(frozen=True) +class _BatchLayout: + counts: np.ndarray + num_examples: int + n_patches: int + pixels_per_patch: int + pooled_per_image: np.ndarray + crops_per_example: np.ndarray + pooled_per_example: np.ndarray + patches_per_image: np.ndarray + + +def _batch_layout( + counts: np.ndarray, + pixel_values: np.ndarray, + image_grids: np.ndarray, + image_num_crops: np.ndarray, +) -> _BatchLayout: + num_examples = counts.shape[0] + _, n_patches, pixels_per_patch = pixel_values.shape + grids = np.asarray(image_grids) + pooled_per_image = (grids[:, 0] * grids[:, 1] + grids[:, 2] * grids[:, 3]).astype(np.int64) + example_for_image = np.repeat(np.arange(num_examples, dtype=np.int64), counts).astype(np.int64) + crops_per_example = np.zeros(num_examples, dtype=np.int64) + np.add.at(crops_per_example, example_for_image, image_num_crops.astype(np.int64)) + pooled_per_example = np.zeros(num_examples, dtype=np.int64) + np.add.at(pooled_per_example, example_for_image, pooled_per_image) + return _BatchLayout( + counts=counts, + num_examples=num_examples, + n_patches=n_patches, + pixels_per_patch=pixels_per_patch, + pooled_per_image=pooled_per_image, + crops_per_example=crops_per_example, + pooled_per_example=pooled_per_example, + patches_per_image=image_num_crops.astype(np.int64) * n_patches, + ) + + +def _allocate_batched_outputs( + layout: _BatchLayout, pixel_values: np.ndarray, image_token_pooling: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + max_crops = int(layout.crops_per_example.max()) if layout.num_examples > 0 else 0 + images = np.full( + (layout.num_examples, max_crops, layout.n_patches, layout.pixels_per_patch), + -1.0, + dtype=pixel_values.dtype, + ) + max_pooled = int(layout.pooled_per_example.max()) if layout.num_examples > 0 else 0 + token_pooling = np.full( + (layout.num_examples, max_pooled, image_token_pooling.shape[-1]), + -1, + dtype=image_token_pooling.dtype, + ) + return images, token_pooling + + +def _offset_example_pooling( + layout: _BatchLayout, + image_token_pooling: np.ndarray, + *, + example_idx: int, + image_offset: int, + pooled_offset: int, +) -> np.ndarray: + example_pooling = image_token_pooling[ + pooled_offset : pooled_offset + int(layout.pooled_per_example[example_idx]) + ].copy() + patch_offset = 0 + row = 0 + for local_image in range(int(layout.counts[example_idx])): + num_pooled = int(layout.pooled_per_image[image_offset + local_image]) + block = example_pooling[row : row + num_pooled] + example_pooling[row : row + num_pooled] = np.where(block >= 0, block + patch_offset, block) + patch_offset += int(layout.patches_per_image[image_offset + local_image]) + row += num_pooled + return example_pooling + + +def build_batched_images( + config: MolmoAct2InputConfig, + input_ids: np.ndarray, + pixel_values: np.ndarray, + image_token_pooling: np.ndarray, + image_grids: np.ndarray, + image_num_crops: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Regroup per-image crops/pooling into per-example padded tensors. + + Mirrors the PyTorch host-side reassembly: infers the image-to-example + mapping from ``image_end`` tokens and offsets pooling indices into each + example's stacked crop patches. + + Returns: + ``(images, token_pooling)`` of shapes ``(N, max_crops, n_patches, pixels)`` + and ``(N, max_pooled, pool_area)``. + + Raises: + ValueError: If image counts cannot be inferred from image-end tokens. + """ + raw_counts = (input_ids == int(config.image_end_token_id)).sum(1) + num_images = int(image_grids.shape[0]) + total_end_tokens = int(raw_counts.sum()) + if num_images == 0: + counts = np.zeros_like(raw_counts) + elif total_end_tokens == num_images: + counts = raw_counts + elif total_end_tokens == 2 * num_images: + counts = raw_counts // 2 + else: + msg = ( + "Could not infer image counts from image-end tokens: " + f"end_tokens={total_end_tokens}, image_grids={num_images}." + ) + raise ValueError(msg) + + layout = _batch_layout(counts, pixel_values, image_grids, image_num_crops) + images, token_pooling = _allocate_batched_outputs(layout, pixel_values, image_token_pooling) + + crop_offset = 0 + pooled_offset = 0 + image_offset = 0 + for example_idx in range(layout.num_examples): + num_example_images = int(layout.counts[example_idx]) + num_example_crops = int(layout.crops_per_example[example_idx]) + images[example_idx, :num_example_crops] = pixel_values[crop_offset : crop_offset + num_example_crops] + + example_pooling = _offset_example_pooling( + layout, + image_token_pooling, + example_idx=example_idx, + image_offset=image_offset, + pooled_offset=pooled_offset, + ) + token_pooling[example_idx, : example_pooling.shape[0]] = example_pooling + + crop_offset += num_example_crops + pooled_offset += int(layout.pooled_per_example[example_idx]) + image_offset += num_example_images + + return images, token_pooling + + +def default_action_dim_is_pad(config: MolmoAct2InputConfig, *, batch_size: int) -> np.ndarray: + """Mark action dimensions beyond the environment action dim as padding. + + Returns: + A boolean mask with padding dimensions marked true. + """ + action_dim_is_pad = np.ones((batch_size, int(config.max_action_dim)), dtype=bool) + if int(config.env_action_dim) > 0: + action_dim_is_pad[:, : int(config.env_action_dim)] = False + return action_dim_is_pad + + +@dataclass(eq=False, repr=False, kw_only=True) +class MolmoAct2ModelInputs(Preprocessor): + """Assemble tokenized prompts and packed images into model inputs.""" + + max_action_dim: int + action_dim: int + bos_token_id: int + pad_token_id: int + image_placeholder_token_id: int + image_start_token_id: int + image_end_token_id: int + image_patch_id: int + image_col_id: int | None + low_res_image_start_token_id: int | None + frame_start_token_id: int | None = None + frame_end_token_id: int | None = None + image_low_res_id: int | None = None + image_size: tuple[int, int] = (378, 378) + patch_size: int = 14 + pooling_size: tuple[int, int] = (2, 2) + image_mean: list[float] | None = None + image_std: list[float] | None = None + image_crop_mode: str = "resize" + image_use_col_tokens: bool = True + use_single_crop_col_tokens: bool | None = False + use_single_crop_start_token: bool = True + image_token_ids: list[int] | None = None + + def __post_init__(self) -> None: + """Build the input layout and image processor.""" + self._layout = MolmoAct2InputConfig( + pad_token_id=self.pad_token_id, + image_placeholder_token_id=self.image_placeholder_token_id, + image_patch_id=self.image_patch_id, + image_start_token_id=self.image_start_token_id, + image_end_token_id=self.image_end_token_id, + image_col_id=self.image_col_id, + low_res_image_start_token_id=self.low_res_image_start_token_id, + frame_start_token_id=self.frame_start_token_id, + frame_end_token_id=self.frame_end_token_id, + image_low_res_id=self.image_low_res_id, + image_use_col_tokens=self.image_use_col_tokens, + use_single_crop_col_tokens=self.use_single_crop_col_tokens, + use_single_crop_start_token=self.use_single_crop_start_token, + max_action_dim=self.max_action_dim, + env_action_dim=self.action_dim, + image_token_ids=self.image_token_ids, + ) + self._image_processor = MolmoAct2ImageProcessor( + crop_mode=self.image_crop_mode, + size={"height": self.image_size[0], "width": self.image_size[1]}, + patch_size=self.patch_size, + pooling_size=self.pooling_size, + image_mean=self.image_mean, + image_std=self.image_std, + ) + + @override + def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: + """Convert tokenized prompts and packed images to graph inputs. + + Returns: + Backbone-ready NumPy arrays. + + Raises: + ValueError: If packed images do not have the expected layout. + """ + input_ids = np.asarray(inputs[TOKENIZED_PROMPT], dtype=np.int64) + attention_mask = np.asarray(inputs[TOKENIZED_PROMPT_MASK], dtype=np.int64) + input_ids, attention_mask = self._insert_bos(input_ids, attention_mask) + + images = np.asarray(inputs[IMAGES], dtype=np.float32) + if images.ndim != _PACKED_IMAGE_NDIM: + msg = f"Expected packed images [N, B, C, H, W], got {images.shape}." + raise ValueError(msg) + num_images, batch_size, channels, height, width = images.shape + flat_images = images.transpose(1, 0, 2, 3, 4).reshape( + batch_size * num_images, + channels, + height, + width, + ) + image_output = self._image_processor(flat_images) + input_ids, attention_mask, token_type_ids = expand_image_placeholders( + config=self._layout, + input_ids=input_ids, + attention_mask=attention_mask, + image_grids=image_output["image_grids"], + ) + batched_images, pooling = build_batched_images( + self._layout, + input_ids, + image_output["pixel_values"], + image_output["image_token_pooling"], + image_output["image_grids"], + image_output["image_num_crops"], + ) + outputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + IMAGES: batched_images.astype(np.float32), + "token_pooling": pooling.astype(np.int64), + "action_dim_is_pad": default_action_dim_is_pad(self._layout, batch_size=batch_size), + } + if token_type_ids is not None: + outputs["token_type_ids"] = token_type_ids + return outputs + + def _insert_bos(self, ids: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + rows = [row_ids[row_mask.astype(bool)] for row_ids, row_mask in zip(ids, mask, strict=True)] + if all(row.size > 0 and int(row[0]) == self.bos_token_id for row in rows): + return ids, mask + rows = [ + row + if row.size > 0 and int(row[0]) == self.bos_token_id + else np.concatenate((np.asarray([self.bos_token_id], dtype=ids.dtype), row)) + for row in rows + ] + width = ids.shape[1] + 1 + output_ids = np.full((len(rows), width), self.pad_token_id, dtype=ids.dtype) + output_mask = np.zeros((len(rows), width), dtype=mask.dtype) + for index, row in enumerate(rows): + output_ids[index, : row.size] = row + output_mask[index, : row.size] = 1 + return output_ids, output_mask + + +__all__ = [ + "MolmoAct2InputConfig", + "MolmoAct2ModelInputs", + "build_batched_images", + "default_action_dim_is_pad", + "expand_image_placeholders", +] diff --git a/src/physicalai/inference/preprocessors/molmoact2/processor.py b/src/physicalai/inference/preprocessors/molmoact2/processor.py new file mode 100644 index 00000000..30cb1e6d --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2/processor.py @@ -0,0 +1,254 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy observation preprocessing for MolmoAct2 inference.""" + +from __future__ import annotations + +import re +from typing import Any + +import cv2 +import numpy as np +from typing_extensions import override + +from physicalai.inference.constants import IMAGES, STATE, TASK +from physicalai.inference.preprocessors.base import Preprocessor +from physicalai.inference.preprocessors.stats_normalizer import StatsNormalizer + +_TRAILING_PUNCTUATION = ".,!?;:" +_IMAGE_NDIM = 4 +_RGB_CHANNELS = 3 +_PREFIX_PATTERNS = tuple( + re.compile(pattern, flags=re.IGNORECASE) + for pattern in ( + r"^(?:task|instruction|language[_ ]instruction|goal)\s*[:\-]\s*", + r"^(?:the\s+task\s+is\s+to|your\s+task\s+is\s+to)\s+", + ) +) + + +def _normalize_text(text: str) -> str: + text = re.sub(r"\s+", " ", str(text or "")).strip() + for pattern in _PREFIX_PATTERNS: + text = pattern.sub("", text, count=1).strip() + return text.rstrip(_TRAILING_PUNCTUATION).strip().lower() + + +def normalization_stats(stats: dict[str, Any]) -> dict[str, np.ndarray]: + """Convert manifest statistics to Studio buffer dtypes. + + Returns: + Float32 statistics and a boolean normalization mask. + """ + return { + name: np.asarray(value, dtype=np.bool_ if name == "mask" else np.float32) + for name, value in stats.items() + if value is not None + } + + +def _wrap(value: str, start: str, end: str, *, enabled: bool) -> str: + if not value or not enabled or (value.startswith(start) and value.endswith(end)): + return value + return f"{start}{value}{end}" + + +def _discrete_state(state: np.ndarray, num_tokens: int) -> str: + if num_tokens <= 0: + msg = f"num_state_tokens must be > 0, got {num_tokens}." + raise ValueError(msg) + state = np.nan_to_num(np.asarray(state, dtype=np.float32), nan=0.0, posinf=1.0, neginf=-1.0) + token_ids = np.rint((np.clip(state, -1.0, 1.0) + 1.0) / 2.0 * (num_tokens - 1)).astype(np.int64) + payload = "".join(f"" for token in token_ids.reshape(-1)) + return f"{payload}" + + +def _build_prompt( + task: str, + state: np.ndarray, + *, + num_state_tokens: int, + setup_type: str, + control_mode: str, + add_setup_tokens: bool, + add_control_tokens: bool, + num_images: int, +) -> str: + setup = _wrap(setup_type, "", "", enabled=add_setup_tokens) + control = _wrap(control_mode, "", "", enabled=add_control_tokens) + prompt = ( + f"The task is to {task}. The setup is {setup}. " + f"The current state of the robot is {_discrete_state(state, num_state_tokens)}. " + f"The expected control mode is {control}. " + "Given these, what action should the robot take to complete the task?" + ) + if num_images == 1: + image_prefix = "<|image|>" + else: + image_prefix = "".join(f"Image {index + 1}<|image|>" for index in range(num_images)) + return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" + + +class MolmoAct2Preprocessor(Preprocessor): + """Prepare normalized prompts and packed images before tokenization.""" + + def __init__( + self, + *, + image_keys: list[str], + state_stats: dict[str, Any] | None = None, + normalization_mode: str = "QUANTILES", + image_size: tuple[int, int] = (378, 378), + num_state_tokens: int = 256, + setup_type: str = "", + control_mode: str = "", + add_setup_tokens: bool = True, + add_control_tokens: bool = True, + ) -> None: + """Store observation preprocessing settings. + + Raises: + ValueError: If no state tokens are available. + """ + if num_state_tokens <= 0: + msg = f"num_state_tokens must be > 0, got {num_state_tokens}." + raise ValueError(msg) + self.image_keys = list(image_keys) + self.image_size = tuple(image_size) + self.num_state_tokens = num_state_tokens + self.setup_type = setup_type + self.control_mode = control_mode + self.add_setup_tokens = add_setup_tokens + self.add_control_tokens = add_control_tokens + self.normalizer = ( + StatsNormalizer( + stats={STATE: normalization_stats(state_stats)}, + mode=normalization_mode.lower(), + features=[STATE], + ) + if state_stats + else None + ) + + @override + def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Prepare one observation batch for tokenizer inference. + + Returns: + Inputs with packed images and encoded prompt text. + + Raises: + ValueError: If a required state, task, or image input is invalid. + """ + outputs = dict(inputs) + state = outputs.get(STATE, outputs.get(f"observation.{STATE}")) + if state is None: + msg = f"MolmoAct2 requires {STATE!r} in its input" + raise ValueError(msg) + state = np.asarray(state, dtype=np.float32) + if state.ndim == 1: + state = state[None, :] + if self.normalizer is not None: + state = self.normalizer({STATE: state})[STATE] + state = np.clip(state, -1.0, 1.0) + + images = self._images(outputs, batch_size=state.shape[0]) + tasks = self._tasks(outputs, batch_size=state.shape[0]) + outputs[IMAGES] = np.stack([self._resize(image) for image in images]) + outputs[TASK] = [ + _build_prompt( + tasks[index], + state[index], + num_state_tokens=self.num_state_tokens, + setup_type=self.setup_type, + control_mode=self.control_mode, + add_setup_tokens=self.add_setup_tokens, + add_control_tokens=self.add_control_tokens, + num_images=len(images), + ) + for index in range(state.shape[0]) + ] + outputs.pop(STATE, None) + outputs.pop(f"observation.{STATE}", None) + return outputs + + def _images(self, inputs: dict[str, Any], *, batch_size: int) -> list[np.ndarray]: + images = self._raw_images(inputs) + output = [] + for image in images: + if image.ndim != _IMAGE_NDIM: + msg = f"Expected BCHW or BHWC image with 3 channels, got {image.shape}" + raise ValueError(msg) + if image.shape[1] == _RGB_CHANNELS: + canonical = image + elif image.shape[-1] == _RGB_CHANNELS: + canonical = image.transpose(0, 3, 1, 2) + else: + msg = f"Expected BCHW or BHWC image with 3 channels, got {image.shape}" + raise ValueError(msg) + if canonical.shape[0] != batch_size: + msg = f"Image batch size mismatch: expected {batch_size}, got {canonical.shape[0]}" + raise ValueError(msg) + output.append(canonical) + return output + + def _raw_images(self, inputs: dict[str, Any]) -> list[np.ndarray]: + container = inputs.get(IMAGES) + images: list[np.ndarray] = [] + for name in self.image_keys: + key = name if name.startswith(f"{IMAGES}.") else f"{IMAGES}.{name}" + if key in inputs: + images.append(np.asarray(inputs[key])) + elif isinstance(container, dict) and name.removeprefix(f"{IMAGES}.") in container: + images.append(np.asarray(container[name.removeprefix(f"{IMAGES}.")])) + if not self.image_keys and isinstance(container, np.ndarray): + images = [container] + elif not self.image_keys and isinstance(container, dict): + keys = sorted(key for key in container if "is_pad" not in str(key)) + images = [np.asarray(container[key]) for key in keys] + elif not self.image_keys and not images: + keys = sorted(key for key in inputs if key.startswith(f"{IMAGES}.") and "is_pad" not in key) + images = [np.asarray(inputs[key]) for key in keys] + if not images: + msg = "MolmoAct2 requires at least one image input" + raise ValueError(msg) + return images + + @staticmethod + def _tasks(inputs: dict[str, Any], *, batch_size: int) -> list[str]: + source = inputs.get(TASK, inputs.get(f"observation.{TASK}", inputs.get("observation.language"))) + if source is None: + msg = f"MolmoAct2 requires {TASK!r} in its input" + raise ValueError(msg) + tasks = [source] * batch_size if isinstance(source, str) else np.asarray(source).reshape(-1).tolist() + if len(tasks) == 1 and batch_size > 1: + tasks *= batch_size + if len(tasks) != batch_size: + msg = f"Expected {batch_size} task strings, got {len(tasks)}" + raise ValueError(msg) + return [_normalize_text(str(task)) for task in tasks] + + def _resize(self, images: np.ndarray) -> np.ndarray: + height, width = self.image_size + output = [] + for image in images: + if image.dtype == np.uint8: + pixels = image + elif np.issubdtype(image.dtype, np.floating): + pixels = image.astype(np.float32) + if float(pixels.max()) <= 1.0: + pixels *= 255.0 + pixels = np.clip(pixels, 0.0, 255.0).astype(np.uint8) + else: + msg = f"Unsupported image dtype: {image.dtype}" + raise ValueError(msg) + resized = cv2.resize(pixels.transpose(1, 2, 0), (width, height), interpolation=cv2.INTER_LINEAR_EXACT) + output.append(resized.transpose(2, 0, 1).astype(np.float32) / 255.0) + return np.stack(output) + + +__all__ = [ + "MolmoAct2Preprocessor", + "normalization_stats", +] diff --git a/src/physicalai/inference/preprocessors/stats_normalizer.py b/src/physicalai/inference/preprocessors/stats_normalizer.py index 61f3a4a1..67067460 100644 --- a/src/physicalai/inference/preprocessors/stats_normalizer.py +++ b/src/physicalai/inference/preprocessors/stats_normalizer.py @@ -187,24 +187,25 @@ def _normalize( Raises: ValueError: If any stats array contains NaN or Inf. """ + transformed = tensor if mode == "mean_std": mean = stats["mean"] std = stats["std"] if not np.all(np.isfinite(mean)) or not np.all(np.isfinite(std)): msg = "mean_std stats contain NaN or Inf — the model artifact may be corrupted" raise ValueError(msg) - return (tensor - mean) / (std + _EPS) + transformed = (tensor - mean) / (std + _EPS) - if mode == "min_max": + elif mode == "min_max": min_val = stats["min"] max_val = stats["max"] if not np.all(np.isfinite(min_val)) or not np.all(np.isfinite(max_val)): msg = "min_max stats contain NaN or Inf — the model artifact may be corrupted" raise ValueError(msg) denom = max_val - min_val + _EPS - return 2.0 * (tensor - min_val) / denom - 1.0 + transformed = 2.0 * (tensor - min_val) / denom - 1.0 - if mode == "quantiles": + elif mode == "quantiles": q01 = stats["q01"] q99 = stats["q99"] if not np.all(np.isfinite(q01)) or not np.all(np.isfinite(q99)): @@ -212,6 +213,7 @@ def _normalize( raise ValueError(msg) denom = q99 - q01 denom = np.where(denom == 0, _EPS, denom) - return 2.0 * (tensor - q01) / denom - 1.0 + transformed = 2.0 * (tensor - q01) / denom - 1.0 - return tensor + mask = stats.get("mask") + return np.where(mask.astype(np.bool_), transformed, tensor) if mask is not None else transformed diff --git a/tests/unit/inference/postprocessors/test_joint_frame.py b/tests/unit/inference/postprocessors/test_joint_frame.py new file mode 100644 index 00000000..c64f6185 --- /dev/null +++ b/tests/unit/inference/postprocessors/test_joint_frame.py @@ -0,0 +1,55 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest + +from physicalai.inference.component_factory import instantiate_component +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.postprocessors import JointFramePostprocessor +from physicalai.inference.postprocessors.joint_frame import JointFrameTransform + + +def test_joint_transform_round_trip_uses_supplied_frame() -> None: + transform = JointFrameTransform(signs=(1.0, -1.0), offsets=(10.0, 20.0)) + robot_values = np.array([[2.0, 3.0, 4.0]], dtype=np.float32) + + checkpoint_values = transform.forward(robot_values) + + np.testing.assert_array_equal(checkpoint_values, [[12.0, 17.0, 4.0]]) + np.testing.assert_array_equal(transform.inverse(checkpoint_values), robot_values) + + +def test_joint_transform_rejects_invalid_frame() -> None: + with pytest.raises(ValueError, match="must match"): + JointFrameTransform(signs=(1.0,), offsets=(0.0, 1.0)) + with pytest.raises(ValueError, match="either -1 or 1"): + JointFrameTransform(signs=(2.0,), offsets=(0.0,)) + + +def test_postprocessor_transforms_configured_feature() -> None: + processor = instantiate_component( + ComponentSpec( + type="joint_frame_postprocess", + feature="action", + signs=[1.0, -1.0], + offsets=[10.0, 20.0], + ) + ) + outputs = { + "action": np.array([[12.0, 17.0, 4.0]], dtype=np.float32), + "other": np.array([5.0]), + } + + assert isinstance(processor, JointFramePostprocessor) + result = processor(outputs) + np.testing.assert_array_equal(result["action"], [[2.0, 3.0, 4.0]]) + np.testing.assert_array_equal(result["other"], outputs["other"]) + np.testing.assert_array_equal(outputs["action"], [[12.0, 17.0, 4.0]]) + + +def test_postprocessor_rejects_missing_feature() -> None: + processor = JointFramePostprocessor(feature="action", signs=[1.0], offsets=[0.0]) + + with pytest.raises(ValueError, match="expected feature 'action'"): + processor({}) diff --git a/tests/unit/inference/postprocessors/test_molmoact2.py b/tests/unit/inference/postprocessors/test_molmoact2.py new file mode 100644 index 00000000..b92ada40 --- /dev/null +++ b/tests/unit/inference/postprocessors/test_molmoact2.py @@ -0,0 +1,59 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest + +from physicalai.inference.component_factory import instantiate_component +from physicalai.inference.constants import ACTION +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.postprocessors import JointFramePostprocessor, MolmoAct2Postprocessor + + +class TestMolmoAct2Postprocessor: + def test_clamps_and_masked_denormalizes_before_joint_transform(self) -> None: + processor = MolmoAct2Postprocessor( + action_key="actions", + action_stats={ + "q01": [0.0, 0.0, 0.0], + "q99": [2.0, 2.0, 2.0], + "mask": [True, False, True], + }, + ) + joint_transform = JointFramePostprocessor( + feature=ACTION, + signs=[1.0, -1.0], + offsets=[0.0, 2.0], + ) + + result = joint_transform(processor({"actions": np.array([[[2.0, 0.5, -2.0]]], dtype=np.float32)})) + + np.testing.assert_allclose(result[ACTION], [[[2.0, 1.5, 0.0]]]) + assert "actions" not in result + + def test_identity_without_stats(self) -> None: + processor = MolmoAct2Postprocessor(action_key=ACTION) + result = processor({ACTION: np.array([[-0.5, 0.5]], dtype=np.float32)}) + np.testing.assert_array_equal(result[ACTION], [[-0.5, 0.5]]) + + def test_missing_action_raises(self) -> None: + with pytest.raises(ValueError, match="expected action key 'model_actions'"): + MolmoAct2Postprocessor(action_key="model_actions")({"other": np.zeros(1)}) + + def test_registry_alias_uses_manifest_action_key(self) -> None: + processor = instantiate_component( + ComponentSpec(type="molmoact2_postprocess", action_key="model_actions"), + ) + + assert isinstance(processor, MolmoAct2Postprocessor) + result = processor( + { + "model_actions": np.array([[-0.5, 0.5]], dtype=np.float32), + "other": np.ones(1), + }, + ) + np.testing.assert_array_equal(result[ACTION], [[-0.5, 0.5]]) + np.testing.assert_array_equal(result["other"], np.ones(1)) + assert "model_actions" not in result diff --git a/tests/unit/inference/postprocessors/test_stats_denormalizer.py b/tests/unit/inference/postprocessors/test_stats_denormalizer.py index 62232f47..cd2e2cb9 100644 --- a/tests/unit/inference/postprocessors/test_stats_denormalizer.py +++ b/tests/unit/inference/postprocessors/test_stats_denormalizer.py @@ -232,6 +232,23 @@ def test_denormalizes_via_stats_param(self) -> None: expected = (np.array([1.9, 3.8, 5.7]) + np.array([0.1, 0.2, 0.3])) / 2.0 np.testing.assert_allclose(result["action"], expected) + def test_masked_dimensions_only(self) -> None: + denormalizer = StatsDenormalizer( + mode="quantiles", + stats={ + "action": { + "q01": [0.0, 0.0, 0.0], + "q99": [2.0, 2.0, 2.0], + "mask": [True, False, True], + }, + }, + ) + outputs = {"action": np.array([[0.0, 0.5, 1.0]], dtype=np.float32)} + + result = denormalizer(outputs) + + np.testing.assert_allclose(result["action"], [[1.0, 0.5, 2.0]]) + class TestStatsDenormalizerLazyLoading: def test_stats_not_loaded_at_init(self, stats_dir: Path) -> None: diff --git a/tests/unit/inference/preprocessors/test_joint_frame.py b/tests/unit/inference/preprocessors/test_joint_frame.py new file mode 100644 index 00000000..13ccc702 --- /dev/null +++ b/tests/unit/inference/preprocessors/test_joint_frame.py @@ -0,0 +1,45 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest + +from physicalai.inference.component_factory import instantiate_component +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.preprocessors import JointFramePreprocessor + + +def test_preprocessor_transforms_configured_feature() -> None: + processor = instantiate_component( + ComponentSpec( + type="joint_frame_preprocess", + feature="state", + signs=[1.0, -1.0], + offsets=[10.0, 20.0], + ) + ) + inputs = { + "state": np.array([[2.0, 3.0, 4.0]], dtype=np.float32), + "other": np.array([5.0]), + } + + assert isinstance(processor, JointFramePreprocessor) + result = processor(inputs) + np.testing.assert_array_equal(result["state"], [[12.0, 17.0, 4.0]]) + np.testing.assert_array_equal(result["other"], inputs["other"]) + np.testing.assert_array_equal(inputs["state"], [[2.0, 3.0, 4.0]]) + + +def test_preprocessor_accepts_observation_prefixed_feature() -> None: + processor = JointFramePreprocessor(feature="state", signs=[-1.0], offsets=[2.0]) + + result = processor({"observation.state": np.array([[3.0]], dtype=np.float32)}) + + np.testing.assert_array_equal(result["observation.state"], [[-1.0]]) + + +def test_preprocessor_rejects_missing_feature() -> None: + processor = JointFramePreprocessor(feature="state", signs=[1.0], offsets=[0.0]) + + with pytest.raises(ValueError, match="expected feature 'state'"): + processor({}) diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py new file mode 100644 index 00000000..30aa6d27 --- /dev/null +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -0,0 +1,314 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest + +from physicalai.inference.component_factory import instantiate_component +from physicalai.inference.constants import IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.postprocessors import MolmoAct2Postprocessor +from physicalai.inference.preprocessors import JointFramePreprocessor, MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.preprocessors.molmoact2.image import MolmoAct2ImageProcessor +from physicalai.inference.preprocessors.molmoact2.inputs import ( + MolmoAct2InputConfig, + build_batched_images, + expand_image_placeholders, +) + + +def _prepare(**kwargs) -> MolmoAct2Preprocessor: + return MolmoAct2Preprocessor( + image_keys=["top", "wrist"], + image_size=(28, 28), + num_state_tokens=4, + setup_type="tabletop", + control_mode="joint", + **kwargs, + ) + + +def _assemble(**kwargs) -> MolmoAct2ModelInputs: + params = { + "max_action_dim": 4, + "action_dim": 2, + "bos_token_id": 1, + "pad_token_id": 0, + "image_placeholder_token_id": 99, + "image_start_token_id": 10, + "image_end_token_id": 12, + "image_patch_id": 11, + "image_col_id": 13, + "low_res_image_start_token_id": 10, + "image_token_ids": [10, 11, 12, 13], + "image_size": (28, 28), + "patch_size": 14, + "pooling_size": (2, 2), + } + params.update(kwargs) + return MolmoAct2ModelInputs(**params) + + +def _observation(*, image_count: int = 1) -> dict: + values = { + STATE: np.array([[-1.0, 1.0]], dtype=np.float32), + TASK: ["Task: Pick up."], + } + for index, key in enumerate(["top", "wrist"][:image_count]): + values[f"{IMAGES}.{key}"] = np.full((1, 3, 28, 28), index * 255, dtype=np.uint8) + return values + + +class TestMolmoAct2Preprocessor: + def test_builds_prompt_and_packs_ordered_cameras(self) -> None: + processor = _prepare() + result = processor(_observation(image_count=2)) + + assert result[IMAGES].shape == (2, 1, 3, 28, 28) + assert float(result[IMAGES][0].max()) == 0.0 + assert float(result[IMAGES][1].min()) == 1.0 + assert result[TASK][0].startswith("Image 1<|image|>Image 2<|image|>") + assert "The task is to pick up." in result[TASK][0] + assert "" in result[TASK][0] + + def test_nested_image_fallback_is_sorted_but_explicit_order_is_preserved(self) -> None: + images = { + "wrist": np.full((1, 3, 28, 28), 255, dtype=np.uint8), + "top": np.zeros((1, 3, 28, 28), dtype=np.uint8), + } + inputs = {STATE: np.zeros((1, 2), dtype=np.float32), TASK: "move", IMAGES: images} + + fallback = MolmoAct2Preprocessor(image_keys=[], image_size=(28, 28))(inputs)[IMAGES] + explicit = MolmoAct2Preprocessor(image_keys=["wrist", "top"], image_size=(28, 28))(inputs)[IMAGES] + + assert float(fallback[0].max()) == 0.0 + assert float(fallback[1].min()) == 1.0 + assert float(explicit[0].min()) == 1.0 + assert float(explicit[1].max()) == 0.0 + + def test_accepts_batched_channels_last_camera_frames(self) -> None: + observation = _observation(image_count=2) + observation[f"{IMAGES}.top"] = np.zeros((1, 28, 28, 3), dtype=np.uint8) + observation[f"{IMAGES}.wrist"] = np.full((1, 28, 28, 3), 255, dtype=np.uint8) + + result = _prepare()(observation) + + assert result[IMAGES].shape == (2, 1, 3, 28, 28) + assert float(result[IMAGES][0].max()) == 0.0 + assert float(result[IMAGES][1].min()) == 1.0 + + def test_applies_masked_normalization_after_joint_transform(self) -> None: + joint_transform = JointFramePreprocessor( + feature=STATE, + signs=[1.0, -1.0], + offsets=[0.0, 2.0], + ) + processor = MolmoAct2Preprocessor( + image_keys=[], + image_size=(28, 28), + state_stats={"q01": [0.0, 0.0], "q99": [2.0, 2.0], "mask": [True, False]}, + ) + inputs = { + STATE: np.array([[1.0, 1.0]], dtype=np.float32), + TASK: "move", + IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8), + } + + result = processor(joint_transform(inputs)) + + assert "" in result[TASK][0] + + def test_supports_mean_std_normalization(self) -> None: + processor = MolmoAct2Preprocessor( + image_keys=[], + image_size=(28, 28), + num_state_tokens=4, + state_stats={"mean": [0.0], "std": [2.0]}, + normalization_mode="MEAN_STD", + ) + result = processor({ + STATE: np.array([[2.0]], dtype=np.float32), + TASK: "move", + IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8), + }) + + assert "" in result[TASK][0] + + def test_preserves_masked_tokenizer_padding(self) -> None: + result = _assemble()({ + TOKENIZED_PROMPT: np.array([[99, 5, 0, 0]], dtype=np.int64), + TOKENIZED_PROMPT_MASK: np.array([[1, 1, 0, 0]], dtype=np.bool_), + IMAGES: np.zeros((1, 1, 3, 28, 28), dtype=np.float32), + }) + + assert result["input_ids"].shape == (1, 7) + assert result["attention_mask"].shape == (1, 7) + assert int(result["attention_mask"].sum()) == 5 + + def test_placeholder_expansion_uses_configured_padding(self) -> None: + config = MolmoAct2InputConfig( + pad_token_id=7, + image_placeholder_token_id=99, + image_patch_id=11, + image_start_token_id=10, + image_end_token_id=12, + ) + + input_ids, attention_mask, _ = expand_image_placeholders( + config=config, + input_ids=np.array([[99, 5], [99, 99]], dtype=np.int64), + attention_mask=np.ones((2, 2), dtype=np.int64), + image_grids=np.array([[1, 1, 0, 0]] * 3, dtype=np.int64), + ) + + assert input_ids[0].tolist() == [10, 11, 12, 5, 7, 7] + assert attention_mask[0].tolist() == [1, 1, 1, 1, 0, 0] + + def test_build_batched_images_supports_multi_crop_grids(self) -> None: + config = MolmoAct2InputConfig( + pad_token_id=0, + image_placeholder_token_id=99, + image_patch_id=11, + image_start_token_id=10, + image_end_token_id=12, + ) + + images, pooling = build_batched_images( + config, + input_ids=np.array([[10, 11, 12, 10, 11, 12], [10, 11, 12, 10, 11, 12]]), + pixel_values=np.arange(8, dtype=np.float32).reshape(2, 4, 1), + image_token_pooling=np.arange(10, dtype=np.int64).reshape(10, 1) % 4, + image_grids=np.array([[1, 1, 2, 2], [1, 1, 2, 2]], dtype=np.int64), + image_num_crops=np.ones(2, dtype=np.int64), + ) + + assert images.shape == (2, 1, 4, 1) + assert pooling.shape == (2, 5, 1) + + def test_rejects_placeholder_image_mismatch(self) -> None: + with pytest.raises(ValueError, match="placeholders"): + _assemble()({ + TOKENIZED_PROMPT: np.array([[99, 99, 5]], dtype=np.int64), + TOKENIZED_PROMPT_MASK: np.ones((1, 3), dtype=np.bool_), + IMAGES: np.zeros((1, 1, 3, 28, 28), dtype=np.float32), + }) + + def test_nullable_single_crop_columns_follow_image_setting(self) -> None: + result = _assemble(use_single_crop_col_tokens=None)({ + TOKENIZED_PROMPT: np.array([[99]], dtype=np.int64), + TOKENIZED_PROMPT_MASK: np.array([[1]], dtype=np.bool_), + IMAGES: np.zeros((1, 1, 3, 28, 28), dtype=np.float32), + }) + + assert result["input_ids"].tolist() == [[1, 10, 11, 13, 12]] + + def test_rejects_missing_state(self) -> None: + with pytest.raises(ValueError, match="state"): + _prepare()({TASK: ["move"], IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8)}) + + def test_model_inputs_is_distinct_class(self) -> None: + assert MolmoAct2ModelInputs is not MolmoAct2Preprocessor + + +class TestMolmoAct2ImageProcessor: + def test_resize_mode_matches_patch_layout(self) -> None: + processor = MolmoAct2ImageProcessor( + crop_mode="resize", + size={"height": 28, "width": 28}, + patch_size=14, + pooling_size=(2, 2), + ) + + result = processor(np.zeros((2, 3, 28, 28), dtype=np.float32)) + + assert result["pixel_values"].shape == (2, 4, 588) + assert result["image_token_pooling"].tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert result["image_grids"].tolist() == [[1, 1, 0, 0], [1, 1, 0, 0]] + + +class TestMolmoAct2ManifestPipeline: + def test_processes_observation_and_action(self, monkeypatch) -> None: + from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer + + class TransformersTokenizer: + name_or_path = "allenai/MolmoAct2" + config = type("Config", (), {"revision": "1dbc166cf8765166998eff31ade2eb64c8a40076"})() + + def __call__(self, tasks, **kwargs): + del tasks, kwargs + return { + "input_ids": np.array([[154629, 7, 0, 0]], dtype=np.int64), + "attention_mask": np.array([[1, 1, 0, 0]], dtype=np.int64), + } + + monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", lambda *args, **kwargs: TransformersTokenizer()) + specs = [ + ComponentSpec( + type="molmoact2", + image_keys=["top"], + state_stats={"q01": [-1.0, -1.0], "q99": [1.0, 1.0]}, + image_size=(28, 28), + ), + ComponentSpec( + type="hf_tokenizer", + tokenizer_name="allenai/MolmoAct2", + revision="1dbc166cf8765166998eff31ade2eb64c8a40076", + max_token_len=4, + ), + ComponentSpec( + type="molmoact2_inputs", + max_action_dim=4, + action_dim=2, + bos_token_id=1, + pad_token_id=0, + image_placeholder_token_id=154629, + image_start_token_id=154624, + image_end_token_id=154625, + image_patch_id=154626, + image_col_id=154627, + low_res_image_start_token_id=154628, + frame_start_token_id=154631, + frame_end_token_id=154632, + image_low_res_id=154630, + image_size=(28, 28), + patch_size=14, + pooling_size=(2, 2), + image_token_ids=[154624, 154625, 154626, 154627, 154628], + ), + ] + values = { + "state": np.array([[0.0, 0.5]], dtype=np.float32), + "task": ["pick up the block"], + "images.top": np.zeros((1, 3, 28, 28), dtype=np.uint8), + } + preprocessor = instantiate_component(specs[0]) + assert isinstance(preprocessor, MolmoAct2Preprocessor) + values = preprocessor(values) + + tokenizer = instantiate_component(specs[1]) + assert isinstance(tokenizer, HFTokenizer) + values = tokenizer(values) + + model_inputs = instantiate_component(specs[2]) + assert isinstance(model_inputs, MolmoAct2ModelInputs) + assert model_inputs._layout.frame_start_token_id == 154631 + assert model_inputs._layout.frame_end_token_id == 154632 + assert model_inputs._layout.image_low_res_id == 154630 + values = model_inputs(values) + + assert set(values) == {"input_ids", "attention_mask", "images", "token_pooling", "action_dim_is_pad", "token_type_ids"} + assert values["images"].shape == (1, 1, 4, 588) + assert values["action_dim_is_pad"].tolist() == [[False, False, True, True]] + + postprocessor = instantiate_component( + ComponentSpec( + type="molmoact2_postprocess", + action_key="action", + action_stats={"q01": [0.0, 0.0], "q99": [2.0, 2.0]}, + ), + ) + assert isinstance(postprocessor, MolmoAct2Postprocessor) + result = postprocessor({"action": np.array([[[0.0, 1.0]]], dtype=np.float32)}) + np.testing.assert_allclose(result["action"], np.array([[[1.0, 2.0]]], dtype=np.float32)) diff --git a/tests/unit/inference/preprocessors/test_stats_normalizer.py b/tests/unit/inference/preprocessors/test_stats_normalizer.py index 440b21b8..4db666d8 100644 --- a/tests/unit/inference/preprocessors/test_stats_normalizer.py +++ b/tests/unit/inference/preprocessors/test_stats_normalizer.py @@ -168,6 +168,23 @@ def test_boundary_values(self, stats_dir: Path) -> None: result_max = normalizer({"observation.state": np.array([1.9, 3.8])}) np.testing.assert_allclose(result_max["observation.state"], np.array([1.0, 1.0])) + def test_masked_dimensions_only(self) -> None: + normalizer = StatsNormalizer( + mode="quantiles", + stats={ + "state": { + "q01": [0.0, 0.0, 0.0], + "q99": [2.0, 2.0, 2.0], + "mask": [True, False, True], + }, + }, + ) + inputs = {"state": np.array([[0.5, 4.0, -0.5]], dtype=np.float32)} + + result = normalizer(inputs) + + np.testing.assert_allclose(result["state"], [[-0.5, 4.0, -1.5]]) + class TestStatsNormalizerIdentity: def test_identity_mode_passthrough(self, stats_dir: Path) -> None: