From 659cabe8209bd428be6fa5b37a7b6b78dc3ab0b1 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:55:18 +0200 Subject: [PATCH 01/22] wip: working inference processor --- src/physicalai/inference/component_factory.py | 2 + src/physicalai/inference/model.py | 23 +- .../inference/postprocessors/__init__.py | 2 + .../inference/postprocessors/molmoact2.py | 81 +++++ .../inference/preprocessors/__init__.py | 2 + .../inference/preprocessors/molmoact2.py | 343 ++++++++++++++++++ 6 files changed, 452 insertions(+), 1 deletion(-) create mode 100644 src/physicalai/inference/postprocessors/molmoact2.py create mode 100644 src/physicalai/inference/preprocessors/molmoact2.py diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index 8d3f0d3d..15bbe82e 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -108,10 +108,12 @@ def __repr__(self) -> str: component_registry.register("hf_tokenizer", "physicalai.inference.preprocessors.HFTokenizer") component_registry.register("ov_tokenizer", "physicalai.inference.preprocessors.OVTokenizer") component_registry.register("pi05", "physicalai.inference.preprocessors.Pi05Preprocessor") +component_registry.register("molmoact2_pre", "physicalai.inference.preprocessors.MolmoAct2Preprocessor") # Postprocessors component_registry.register("denormalize", "physicalai.inference.postprocessors.StatsDenormalizer") component_registry.register("action_chunk_trimmer", "physicalai.inference.postprocessors.ActionChunkTrimmer") +component_registry.register("molmoact2_post", "physicalai.inference.postprocessors.MolmoAct2Postprocessor") def resolve_artifact(spec: ComponentSpec, export_dir: Path) -> ComponentSpec: diff --git a/src/physicalai/inference/model.py b/src/physicalai/inference/model.py index e48d2b2f..5a4b26ba 100644 --- a/src/physicalai/inference/model.py +++ b/src/physicalai/inference/model.py @@ -273,9 +273,30 @@ def select_action(self, observation: dict[str, np.ndarray]) -> np.ndarray: >>> next_obs, reward, done = env.step(action) """ if not self._action_buffer: - self._action_buffer.extend(self.predict_action_chunk(observation)) + action_chunk = self.predict_action_chunk(observation) + self._action_buffer.extend(action_chunk[: self._effective_chunk_size()]) return self._action_buffer.popleft() + def _effective_chunk_size(self) -> int: + """Return the number of actions to queue per model invocation. + + Preference order: + 1. Runner-declared ``chunk_size`` from the manifest. + 2. ACTION output feature leading dimension (when declared as ``(T, D)``). + 3. Fallback to 1. + """ + runner_chunk = int(self.chunk_size) + if runner_chunk > 1: + return runner_chunk + + for feature in self.output_features: + if feature.name == ACTION and len(feature.shape) >= 2: + action_chunk = int(feature.shape[0]) + if action_chunk > 0: + return action_chunk + + return 1 + def predict_action_chunk(self, observation: dict[str, np.ndarray]) -> np.ndarray: """Predict a chunk of actions for the given observation. diff --git a/src/physicalai/inference/postprocessors/__init__.py b/src/physicalai/inference/postprocessors/__init__.py index 64a45d42..e4e52438 100644 --- a/src/physicalai/inference/postprocessors/__init__.py +++ b/src/physicalai/inference/postprocessors/__init__.py @@ -9,11 +9,13 @@ 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.molmoact2 import MolmoAct2Postprocessor from physicalai.inference.postprocessors.stats_denormalizer import StatsDenormalizer __all__ = [ "ActionChunkTrimmer", "ActionNormalizer", + "MolmoAct2Postprocessor", "Postprocessor", "StatsDenormalizer", ] diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py new file mode 100644 index 00000000..ff31c288 --- /dev/null +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -0,0 +1,81 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""MolmoAct2 inference postprocessor.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from physicalai.inference.constants import ACTION + +from .base import Postprocessor + +_EPS = 1e-8 + + +class MolmoAct2Postprocessor(Postprocessor): + """Map model outputs to ``action`` and restore action-space scaling.""" + + def __init__( + self, + *, + action_key: str | None = None, + env_action_dim: int | None = None, + action_stats: dict[str, list[float] | np.ndarray] | None = None, + ) -> None: + self._action_key = action_key + self._env_action_dim = int(env_action_dim) if env_action_dim is not None else None + + self._q01: np.ndarray | None = None + self._q99: np.ndarray | None = None + self._mask: np.ndarray | None = None + if action_stats is not None: + q01 = action_stats.get("q01") + q99 = action_stats.get("q99") + if q01 is not None and q99 is not None: + self._q01 = np.asarray(q01, dtype=np.float32) + self._q99 = np.asarray(q99, dtype=np.float32) + mask = action_stats.get("mask") + if mask is not None: + self._mask = np.asarray(mask, dtype=bool) + + def _resolve_action_key(self, outputs: dict[str, np.ndarray]) -> str: + if ACTION in outputs: + return ACTION + if self._action_key is not None: + return self._action_key + if "actions" in outputs: + return "actions" + return next(iter(outputs)) + + def _denormalize(self, action: np.ndarray) -> np.ndarray: + if self._q01 is None or self._q99 is None: + return action + denom = self._q99 - self._q01 + denom = np.where(denom == 0, _EPS, denom) + denorm = (action + 1.0) * denom / 2.0 + self._q01 + if self._mask is not None: + mask = self._mask + while mask.ndim < denorm.ndim: + mask = np.expand_dims(mask, axis=0) + denorm = np.where(mask, denorm, action) + return denorm + + def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + result = dict(outputs) + action_key = self._resolve_action_key(result) + action = np.asarray(result.pop(action_key)) + + if self._env_action_dim is not None: + action = action[..., : self._env_action_dim] + + action = np.clip(action, -1.0, 1.0) + action = self._denormalize(action) + result[ACTION] = action.astype(np.float32) + return result + + +__all__ = ["MolmoAct2Postprocessor"] diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index c8db365d..02b22bf2 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -10,6 +10,7 @@ from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer from physicalai.inference.preprocessors.lambda_processor import LambdaPreprocessor +from physicalai.inference.preprocessors.molmoact2 import MolmoAct2Preprocessor from physicalai.inference.preprocessors.new_line import NewLinePreprocessor from physicalai.inference.preprocessors.ov_tokenizer import OVTokenizer from physicalai.inference.preprocessors.pi05 import Pi05Preprocessor @@ -20,6 +21,7 @@ __all__ = [ "HFTokenizer", "LambdaPreprocessor", + "MolmoAct2Preprocessor", "NewLinePreprocessor", "OVTokenizer", "Pi05Preprocessor", diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py new file mode 100644 index 00000000..e76ac3c6 --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -0,0 +1,343 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""MolmoAct2 inference preprocessor. + +Builds MolmoAct2 prompts from raw task/state/images, tokenizes text, inserts BOS, +and emits model-facing tensors that are expected outside the exported model graph. +""" + +from __future__ import annotations + +import re +from typing import Any + +import numpy as np + +from physicalai.inference.constants import IMAGES, STATE, TASK + +from .base import Preprocessor + +ACTION_OUTPUT_TOKEN = "" +SETUP_START_TOKEN = "" +SETUP_END_TOKEN = "" +CONTROL_START_TOKEN = "" +CONTROL_END_TOKEN = "" +STATE_START_TOKEN = "" +STATE_END_TOKEN = "" +STATE_TOKEN_PREFIX = " str: + normalized = re.sub(r"\s+", " ", str(text or "")).strip() + if not normalized: + return "" + for pattern in _PREFIX_PATTERNS: + normalized = pattern.sub("", normalized, count=1).strip() + normalized = normalized.rstrip(_TRAILING_PUNCTUATION).strip() + return normalized.lower() + + +def _wrap_setup_text(setup_type: str, add_setup_tokens: bool) -> str: + if not setup_type: + return "" + if not add_setup_tokens: + return setup_type + if setup_type.startswith(SETUP_START_TOKEN) and setup_type.endswith(SETUP_END_TOKEN): + return setup_type + return f"{SETUP_START_TOKEN}{setup_type}{SETUP_END_TOKEN}" + + +def _wrap_control_text(control_mode: str, add_control_tokens: bool) -> str: + if not control_mode: + return "" + if not add_control_tokens: + return control_mode + if control_mode.startswith(CONTROL_START_TOKEN) and control_mode.endswith(CONTROL_END_TOKEN): + return control_mode + return f"{CONTROL_START_TOKEN}{control_mode}{CONTROL_END_TOKEN}" + + +def _build_discrete_state_string(state: np.ndarray, num_state_tokens: int) -> str: + if num_state_tokens <= 0: + msg = f"num_state_tokens must be > 0, got {num_state_tokens}." + raise ValueError(msg) + arr = np.asarray(state, dtype=np.float32) + arr = np.nan_to_num(arr, nan=0.0, posinf=1.0, neginf=-1.0) + arr = np.clip(arr, -1.0, 1.0) + scaled = (arr + 1.0) / 2.0 * float(num_state_tokens - 1) + token_ids = np.clip(np.rint(scaled).astype(np.int64), 0, int(num_state_tokens) - 1).reshape(-1) + return f"{STATE_START_TOKEN}{''.join(f'{STATE_TOKEN_PREFIX}{int(token_id)}>' for token_id in token_ids)}{STATE_END_TOKEN}" + + +def _build_robot_text( + *, + task: str, + discrete_state_string: str, + setup_type: str, + control_mode: str, + add_setup_tokens: bool, + add_control_tokens: bool, + num_images: int, +) -> str: + setup_text = _wrap_setup_text(setup_type, add_setup_tokens=add_setup_tokens) + control_text = _wrap_control_text(control_mode, add_control_tokens=add_control_tokens) + state_clause = f" The current state of the robot is {discrete_state_string}." if discrete_state_string else "" + prompt = ( + f"The task is to {task}. The setup is {setup_text}.{state_clause} " + f"The expected control mode is {control_text}. Given these, what action should the robot take to complete the task?" + ) + if num_images <= 0: + image_prefix = "" + elif num_images == 1: + image_prefix = IMAGE_PROMPT + else: + image_prefix = "".join(f"Image {idx + 1}{IMAGE_PROMPT}" for idx in range(num_images)) + return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{ACTION_OUTPUT_TOKEN}" + + +class MolmoAct2Preprocessor(Preprocessor): + """Build MolmoAct2 token and image inputs from raw inference observations.""" + + def __init__( + self, + tokenizer_name_or_path: str, + *, + num_state_tokens: int = 256, + setup_type: str = "", + control_mode: str = "", + add_setup_tokens: bool = False, + add_control_tokens: bool = False, + state_stats: dict[str, list[float] | np.ndarray] | None = None, + image_keys: list[str] | None = None, + ) -> None: + self.tokenizer_name_or_path = tokenizer_name_or_path + self.num_state_tokens = int(num_state_tokens) + self.setup_type = str(setup_type or "") + self.control_mode = str(control_mode or "") + self.add_setup_tokens = bool(add_setup_tokens) + self.add_control_tokens = bool(add_control_tokens) + self.image_keys = list(image_keys or []) + self._tokenizer: Any = None + + self._state_q01: np.ndarray | None = None + self._state_q99: np.ndarray | None = None + self._state_mask: np.ndarray | None = None + if state_stats is not None: + q01 = state_stats.get("q01") + q99 = state_stats.get("q99") + if q01 is not None and q99 is not None: + self._state_q01 = np.asarray(q01, dtype=np.float32) + self._state_q99 = np.asarray(q99, dtype=np.float32) + mask = state_stats.get("mask") + if mask is not None: + self._state_mask = np.asarray(mask, dtype=bool) + + @property + def tokenizer(self) -> Any: + if self._tokenizer is None: + from transformers import Qwen2Tokenizer # noqa: PLC0415 + + self._tokenizer = Qwen2Tokenizer.from_pretrained( + self.tokenizer_name_or_path, + local_files_only=False, + ) + return self._tokenizer + + @staticmethod + def _insert_bos( + input_ids: np.ndarray, + attention_mask: np.ndarray, + bos_token_id: int, + pad_token_id: int, + ) -> tuple[np.ndarray, np.ndarray]: + if input_ids.ndim == 1: + input_ids = input_ids[None, :] + attention_mask = attention_mask[None, :] + squeeze = True + else: + squeeze = False + + batch_size, seq_len = input_ids.shape + if seq_len == 0: + out_ids = np.full((batch_size, 1), bos_token_id, dtype=input_ids.dtype) + out_mask = np.ones((batch_size, 1), dtype=attention_mask.dtype) + return (out_ids[0], out_mask[0]) if squeeze else (out_ids, out_mask) + + first_valid = (attention_mask == 1).argmax(axis=-1) + if np.all(input_ids[np.arange(batch_size), first_valid] == bos_token_id): + return (input_ids[0], attention_mask[0]) if squeeze else (input_ids, attention_mask) + + out_ids = np.full((batch_size, seq_len + 1), pad_token_id, dtype=input_ids.dtype) + out_mask = np.zeros((batch_size, seq_len + 1), dtype=attention_mask.dtype) + + src = np.tile(np.arange(seq_len), (batch_size, 1)) + valid = src >= first_valid[:, None] + tgt = src + 1 + batch_idx = np.tile(np.arange(batch_size)[:, None], (1, seq_len)) + + out_ids[batch_idx[valid], tgt[valid]] = input_ids[valid] + out_mask[batch_idx[valid], tgt[valid]] = 1 + out_ids[np.arange(batch_size), first_valid] = bos_token_id + out_mask[np.arange(batch_size), first_valid] = 1 + return (out_ids[0], out_mask[0]) if squeeze else (out_ids, out_mask) + + def _normalize_state(self, state: np.ndarray) -> np.ndarray: + state = np.asarray(state, dtype=np.float32) + if self._state_q01 is None or self._state_q99 is None: + return np.clip(state, -1.0, 1.0) + + denom = self._state_q99 - self._state_q01 + denom = np.where(denom == 0, _EPS, denom) + normalized = 2.0 * (state - self._state_q01) / denom - 1.0 + if self._state_mask is not None: + mask = self._state_mask + while mask.ndim < normalized.ndim: + mask = np.expand_dims(mask, axis=0) + normalized = np.where(mask, normalized, state) + return np.clip(normalized, -1.0, 1.0) + + def _extract_state(self, inputs: dict[str, Any]) -> np.ndarray: + raw_state = inputs.get(STATE) + if raw_state is None: + raw_state = inputs.get(f"observation.{STATE}") + if raw_state is None: + msg = "MolmoAct2 inference preprocessor requires state." + raise ValueError(msg) + + state = np.asarray(raw_state, dtype=np.float32) + if state.ndim == 1: + state = state[None, :] + return self._normalize_state(state) + + @staticmethod + def _extract_tasks(inputs: dict[str, Any], batch_size: int) -> list[str]: + task_source = inputs.get(TASK) + if task_source is None: + task_source = inputs.get(f"observation.{TASK}") + + if task_source is None: + tasks = [""] * batch_size + elif isinstance(task_source, str): + tasks = [task_source] * batch_size + elif isinstance(task_source, (list, tuple, np.ndarray)): + tasks = [str(item) for item in list(task_source)] + else: + tasks = [str(task_source)] + + if len(tasks) == 1 and batch_size > 1: + tasks = tasks * batch_size + if len(tasks) != batch_size: + msg = f"Expected {batch_size} task strings, got {len(tasks)}." + raise ValueError(msg) + return [_normalize_text(task) for task in tasks] + + def _resolve_image_arrays(self, inputs: dict[str, Any]) -> list[np.ndarray]: + images_value = inputs.get(IMAGES) + if isinstance(images_value, dict): + if self.image_keys: + return [np.asarray(images_value[key]) for key in self.image_keys if key in images_value] + return [np.asarray(value) for value in images_value.values()] + if images_value is not None and not isinstance(images_value, (str, bytes)): + return [np.asarray(images_value)] + + flat_keys: list[str] = [] + if self.image_keys: + flat_keys = [f"{IMAGES}.{key}" for key in self.image_keys if f"{IMAGES}.{key}" in inputs] + if not flat_keys: + flat_keys = [key for key in inputs if str(key).startswith(f"{IMAGES}.") and "is_pad" not in str(key)] + flat_keys.sort() + return [np.asarray(inputs[key]) for key in flat_keys] + + @staticmethod + def _as_bchw_batch(array: np.ndarray) -> np.ndarray: + arr = np.asarray(array) + if arr.ndim == 3: + if int(arr.shape[0]) != 3: + msg = f"Expected CHW image tensor with 3 channels, got shape {arr.shape}" + raise ValueError(msg) + arr = arr[None, ...] + if arr.ndim != 4: + msg = f"Expected BCHW image tensor, got shape {arr.shape}" + raise ValueError(msg) + if int(arr.shape[1]) != 3: + msg = f"Expected BCHW image tensor with 3 channels, got shape {arr.shape}" + raise ValueError(msg) + + if arr.dtype == np.uint8: + arr = arr.astype(np.float32) / 255.0 + else: + arr = arr.astype(np.float32) + return arr + + def _extract_images_by_example(self, inputs: dict[str, Any], batch_size: int) -> list[list[np.ndarray]]: + arrays = self._resolve_image_arrays(inputs) + if not arrays: + msg = "MolmoAct2 inference preprocessor requires at least one image input." + raise ValueError(msg) + + images_by_example: list[list[np.ndarray]] = [[] for _ in range(batch_size)] + for arr in arrays: + bchw = self._as_bchw_batch(arr) + if int(bchw.shape[0]) != batch_size: + msg = f"Image batch size mismatch: expected {batch_size}, got {bchw.shape[0]}" + raise ValueError(msg) + for idx in range(batch_size): + images_by_example[idx].append(bchw[idx]) + return images_by_example + + def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.ndarray]: + inputs_dict = dict(inputs) + + state = self._extract_state(inputs_dict) + batch_size = int(state.shape[0]) + tasks = self._extract_tasks(inputs_dict, batch_size) + images_by_example = self._extract_images_by_example(inputs_dict, batch_size) + + prompt_texts: list[str] = [] + flat_images: list[np.ndarray] = [] + for idx in range(batch_size): + flat_images.extend(images_by_example[idx]) + discrete_state = _build_discrete_state_string(state[idx], self.num_state_tokens) + prompt_texts.append( + _build_robot_text( + task=tasks[idx], + discrete_state_string=discrete_state, + 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_by_example[idx]), + ) + ) + + image_batch = np.stack(flat_images, axis=0).astype(np.float32) if flat_images else np.empty((0, 3, 0, 0), dtype=np.float32) + + text_inputs = self.tokenizer(prompt_texts, padding=True) + input_ids = np.asarray(text_inputs["input_ids"], dtype=np.int64) + attention_mask = np.asarray(text_inputs["attention_mask"], dtype=np.int64) + + bos_token_id = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id + pad_token_id = self.tokenizer.pad_token_id + input_ids, attention_mask = self._insert_bos(input_ids, attention_mask, int(bos_token_id), int(pad_token_id)) + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "image_placeholder_token_id": np.asarray( + int(self.tokenizer.convert_tokens_to_ids(IMAGE_PROMPT)), + dtype=np.int64, + ), + "images_bchw": image_batch, + } From ef670f076046ce3f80d60b3d904f0b914bd9a748 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:26:28 +0200 Subject: [PATCH 02/22] wip: changed preprocessor --- pyproject.toml | 2 +- .../inference/preprocessors/molmoact2.py | 41 ++- .../preprocessors/molmoact2_image.py | 310 ++++++++++++++++++ 3 files changed, 343 insertions(+), 10 deletions(-) create mode 100644 src/physicalai/inference/preprocessors/molmoact2_image.py diff --git a/pyproject.toml b/pyproject.toml index 7a69c92e..cb3d8c7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] dependencies = [ - "numpy>=1.24", + "numpy", "loguru>=0.7", # target PyTurboJPEG 1.x for libjpeg-turbo 2.x compatibility "PyTurboJPEG<2; sys_platform == 'linux'", diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index e76ac3c6..e10064ae 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -14,7 +14,7 @@ import numpy as np -from physicalai.inference.constants import IMAGES, STATE, TASK +from physicalai.inference.constants import IMAGE_MASKS, IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK from .base import Preprocessor @@ -121,6 +121,7 @@ def __init__( add_control_tokens: bool = False, state_stats: dict[str, list[float] | np.ndarray] | None = None, image_keys: list[str] | None = None, + processor_config: dict[str, Any] | None = None, ) -> None: self.tokenizer_name_or_path = tokenizer_name_or_path self.num_state_tokens = int(num_state_tokens) @@ -130,6 +131,7 @@ def __init__( self.add_control_tokens = bool(add_control_tokens) self.image_keys = list(image_keys or []) self._tokenizer: Any = None + del processor_config self._state_q01: np.ndarray | None = None self._state_q99: np.ndarray | None = None @@ -297,6 +299,28 @@ def _extract_images_by_example(self, inputs: dict[str, Any], batch_size: int) -> images_by_example[idx].append(bchw[idx]) return images_by_example + @staticmethod + def _pack_images(images_by_example: list[list[np.ndarray]]) -> tuple[np.ndarray, np.ndarray]: + batch_size = len(images_by_example) + if batch_size == 0: + return np.empty((0, 0, 3, 0, 0), dtype=np.float32), np.empty((0, 0), dtype=bool) + + num_images = len(images_by_example[0]) + if any(len(example_images) != num_images for example_images in images_by_example): + msg = "MolmoAct2 requires a consistent number of images per batch element." + raise ValueError(msg) + + if num_images == 0: + return np.empty((0, batch_size, 3, 0, 0), dtype=np.float32), np.empty((0, batch_size), dtype=bool) + + image_slots: list[np.ndarray] = [] + image_masks: list[np.ndarray] = [] + for image_idx in range(num_images): + slot_images = [images_by_example[batch_idx][image_idx].astype(np.float32, copy=False) for batch_idx in range(batch_size)] + image_slots.append(np.stack(slot_images, axis=0)) + image_masks.append(np.ones((batch_size,), dtype=bool)) + return np.stack(image_slots, axis=0), np.stack(image_masks, axis=0) + def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.ndarray]: inputs_dict = dict(inputs) @@ -322,7 +346,8 @@ def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.nd ) ) - image_batch = np.stack(flat_images, axis=0).astype(np.float32) if flat_images else np.empty((0, 3, 0, 0), dtype=np.float32) + del flat_images + images, image_masks = self._pack_images(images_by_example) text_inputs = self.tokenizer(prompt_texts, padding=True) input_ids = np.asarray(text_inputs["input_ids"], dtype=np.int64) @@ -333,11 +358,9 @@ def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.nd input_ids, attention_mask = self._insert_bos(input_ids, attention_mask, int(bos_token_id), int(pad_token_id)) return { - "input_ids": input_ids, - "attention_mask": attention_mask, - "image_placeholder_token_id": np.asarray( - int(self.tokenizer.convert_tokens_to_ids(IMAGE_PROMPT)), - dtype=np.int64, - ), - "images_bchw": image_batch, + TOKENIZED_PROMPT: input_ids, + TOKENIZED_PROMPT_MASK: attention_mask, + STATE: state, + IMAGES: images, + IMAGE_MASKS: image_masks, } diff --git a/src/physicalai/inference/preprocessors/molmoact2_image.py b/src/physicalai/inference/preprocessors/molmoact2_image.py new file mode 100644 index 00000000..b110313e --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2_image.py @@ -0,0 +1,310 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy image preprocessing helpers for MolmoAct2 inference.""" + +from __future__ import annotations + +import cv2 +import numpy as np + + +def _normalize_image(image: np.ndarray, image_mean: list[float], image_std: list[float]) -> np.ndarray: + if np.allclose(image_mean, [0.5, 0.5, 0.5]) and np.allclose(image_std, [0.5, 0.5, 0.5]): + return image * np.asarray(2.0, dtype=np.float32) - np.asarray(1.0, dtype=np.float32) + image = image.astype(np.float32) + image -= np.asarray(image_mean, dtype=np.float32)[None, None, :] + image /= np.asarray(image_std, dtype=np.float32)[None, None, :] + return image + + +def _resize_image(image: np.ndarray, desired_output_size: list[int]) -> np.ndarray: + height, width = int(desired_output_size[0]), int(desired_output_size[1]) + resized = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR) + if resized.ndim == 2: + resized = resized[:, :, None] + + if np.issubdtype(image.dtype, np.floating): + resized = np.clip(resized, 0.0, 1.0).astype(np.float32) + else: + resized = resized.astype(np.float32) / 255.0 + return resized + + +def _select_tiling(h: int, w: int, patch_size: int, max_num_crops: int) -> np.ndarray: + tilings: list[tuple[int, int]] = [] + for i in range(1, max_num_crops + 1): + for j in range(1, max_num_crops + 1): + if i * j <= max_num_crops: + tilings.append((i, j)) + tilings.sort(key=lambda x: (x[0] * x[1], x[0])) + candidate_tilings = np.asarray(tilings, dtype=np.int32) + candidate_resolutions = candidate_tilings * patch_size + + original_size = np.asarray([h, w], dtype=np.float32) + with np.errstate(divide="ignore"): + required_scale = candidate_resolutions.astype(np.float32) / original_size[None, :] + required_scale = np.min(required_scale, axis=-1, keepdims=True) + if np.all(required_scale < 1): + ix = int(np.argmax(required_scale)) + else: + required_scale = np.where(required_scale < 1.0, 1e10, required_scale) + ix = int(np.argmin(required_scale)) + return candidate_tilings[ix] + + +def _build_resized_image( + image: np.ndarray, + base_image_input_size: list[int], + image_mean: list[float], + image_std: list[float], + image_patch_size: int, +) -> tuple[np.ndarray, np.ndarray]: + resized = _resize_image(image, base_image_input_size) + resized = _normalize_image(resized, image_mean, image_std) + resized = resized[None, ...] + crop_patch_w = base_image_input_size[1] // image_patch_size + crop_patch_h = base_image_input_size[0] // image_patch_size + resize_idx = np.arange(crop_patch_w * crop_patch_h).reshape([crop_patch_h, crop_patch_w]) + return resized, resize_idx + + +def _build_overlapping_crops( + image: np.ndarray, + max_crops: int, + overlap_margins: list[int], + base_image_input_size: list[int], + image_mean: list[float], + image_std: list[float], + image_patch_size: int, +) -> tuple[np.ndarray, np.ndarray]: + left_margin, right_margin = overlap_margins + total_margin_pixels = image_patch_size * (right_margin + left_margin) + crop_patches = base_image_input_size[0] // image_patch_size + crop_window_patches = crop_patches - (right_margin + left_margin) + crop_window_size = crop_window_patches * image_patch_size + crop_patch_w = base_image_input_size[1] // image_patch_size + crop_patch_h = base_image_input_size[0] // image_patch_size + + original_image_h, original_image_w = image.shape[:2] + crop_size = base_image_input_size[0] + + tiling = _select_tiling( + original_image_h - total_margin_pixels, + original_image_w - total_margin_pixels, + crop_window_size, + max_crops, + ) + + src = _resize_image( + image, + [tiling[0] * crop_window_size + total_margin_pixels, tiling[1] * crop_window_size + total_margin_pixels], + ) + src = _normalize_image(src, image_mean, image_std) + + n_crops = int(tiling[0] * tiling[1]) + crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype) + patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32) + + on_crop = 0 + for i in range(int(tiling[0])): + y0 = i * crop_window_size + for j in range(int(tiling[1])): + x0 = j * crop_window_size + crop_arr[on_crop] = src[y0 : y0 + crop_size, x0 : x0 + crop_size] + patch_idx = np.arange(crop_patch_w * crop_patch_h).reshape(crop_patch_h, crop_patch_w) + patch_idx += on_crop * crop_patch_h * crop_patch_w + + if i != 0: + patch_idx[:left_margin, :] = -1 + if j != 0: + patch_idx[:, :left_margin] = -1 + if i != int(tiling[0]) - 1: + patch_idx[-right_margin:, :] = -1 + if j != int(tiling[1]) - 1: + patch_idx[:, -right_margin:] = -1 + patch_idx_arr[on_crop] = patch_idx + on_crop += 1 + + patch_idx_arr = patch_idx_arr.reshape(int(tiling[0]), int(tiling[1]), crop_patch_h, crop_patch_w) + patch_idx_arr = patch_idx_arr.transpose(0, 2, 1, 3).reshape(-1) + patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape( + src.shape[0] // image_patch_size, + src.shape[1] // image_patch_size, + ) + return crop_arr, patch_idx_arr + + +def _batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray: + n_crops, h, w, c = array.shape + h_patches = h // patch_size + w_patches = w // patch_size + array = array.reshape(n_crops, h_patches, patch_size, w_patches, patch_size, c) + array = array.transpose(0, 1, 3, 2, 4, 5) + return array.reshape(n_crops, h_patches * w_patches, patch_size * patch_size * c) + + +def _arange_for_pooling(idx_arr: np.ndarray, pool_h: int, pool_w: int) -> np.ndarray: + h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0] + w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1] + idx_arr = np.pad( + idx_arr, + [[h_pad // 2, (h_pad + 1) // 2], [w_pad // 2, (w_pad + 1) // 2]], + mode="constant", + constant_values=-1, + ) + blocks_h = idx_arr.shape[0] // pool_h + blocks_w = idx_arr.shape[1] // pool_w + idx_arr = idx_arr.reshape(blocks_h, pool_h, blocks_w, pool_w) + idx_arr = idx_arr.transpose(0, 2, 1, 3) + return idx_arr.reshape(blocks_h, blocks_w, pool_h * pool_w) + + +def _image_to_patches_and_grids( + image: np.ndarray, + max_crops: int, + overlap_margins: list[int], + base_image_input_size: list[int], + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + image_pooling_w: int, + image_pooling_h: int, + crop_mode: str, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + crop_patch_w = base_image_input_size[1] // image_patch_size + crop_patch_h = base_image_input_size[0] // image_patch_size + + if crop_mode == "resize": + resized, resize_idx = _build_resized_image( + image, + base_image_input_size, + image_mean, + image_std, + image_patch_size, + ) + resize_idx = _arange_for_pooling(resize_idx, image_pooling_h, image_pooling_w) + resized_h, resized_w = resize_idx.shape[:2] + resize_idx = resize_idx.reshape(-1, image_pooling_h * image_pooling_w) + image_grid = [np.asarray([resized_h, resized_w, 0, 0])] + return np.stack(image_grid, 0), _batch_pixels_to_patches(resized, image_patch_size), resize_idx + + if crop_mode not in {"overlap-and-resize-c2", "overlap-and-resize"}: + msg = f"Unsupported MolmoAct2 image crop_mode {crop_mode!r}." + raise ValueError(msg) + + crop_arr, patch_idx_arr = _build_overlapping_crops( + image, + max_crops, + overlap_margins, + base_image_input_size, + image_mean, + image_std, + image_patch_size, + ) + pooling_idx = _arange_for_pooling(patch_idx_arr, image_pooling_h, image_pooling_w) + h, w = pooling_idx.shape[:2] + pooling_idx = pooling_idx.reshape(-1, image_pooling_h * image_pooling_w) + + resized, resize_idx = _build_resized_image( + image, + base_image_input_size, + image_mean, + image_std, + image_patch_size, + ) + crop_arr = np.concatenate([resized, crop_arr], axis=0) + + resize_idx = _arange_for_pooling(resize_idx, image_pooling_h, image_pooling_w) + resized_h, resized_w = resize_idx.shape[:2] + resize_idx = resize_idx.reshape(-1, image_pooling_h * image_pooling_w) + + pooling_idx = np.where(pooling_idx >= 0, pooling_idx + crop_patch_h * crop_patch_w, -1) + pooling_idx = np.concatenate([resize_idx, pooling_idx], axis=0) + image_grid = [np.asarray([resized_h, resized_w, h, w])] + return np.stack(image_grid, 0), _batch_pixels_to_patches(crop_arr, image_patch_size), pooling_idx + + +def _to_hwc_uint8(images_bchw: np.ndarray) -> list[np.ndarray]: + out: list[np.ndarray] = [] + for image in images_bchw: + img = image + if np.issubdtype(img.dtype, np.floating): + if float(np.max(img)) <= 1.0: + img = img * 255.0 + img = np.clip(img, 0.0, 255.0).astype(np.uint8) + elif img.dtype != np.uint8: + img = np.clip(img, 0, 255).astype(np.uint8) + out.append(np.transpose(img, (1, 2, 0))) + return out + + +class MolmoAct2ImageProcessor: + """NumPy image processor producing MolmoAct2 patch tensors and pooling metadata.""" + + def __init__( + self, + size: dict[str, int] | None = None, + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_convert_rgb: bool = True, + max_crops: int = 8, + overlap_margins: list[int] | None = None, + crop_mode: str = "overlap-and-resize-c2", + patch_size: int = 14, + pooling_size: list[int] | None = None, + ) -> None: + self.size = size if size is not None else {"height": 378, "width": 378} + self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5] + self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5] + self.do_convert_rgb = do_convert_rgb + self.max_crops = int(max_crops) + self.overlap_margins = overlap_margins if overlap_margins is not None else [4, 4] + self.crop_mode = crop_mode + self.patch_size = int(patch_size) + self.pooling_size = pooling_size if pooling_size is not None else [2, 2] + + def __call__(self, images_bchw: np.ndarray) -> dict[str, np.ndarray]: + image_list = _to_hwc_uint8(images_bchw) + patch_batches: list[np.ndarray] = [] + pooling_batches: list[np.ndarray] = [] + grids: list[np.ndarray] = [] + image_num_crops: list[int] = [] + + base_image_input_size = [int(self.size["height"]), int(self.size["width"])] + pool_h, pool_w = int(self.pooling_size[0]), int(self.pooling_size[1]) + + for image in image_list: + image_grid, crops, pooled_idx = _image_to_patches_and_grids( + image, + self.max_crops, + self.overlap_margins, + base_image_input_size, + self.image_mean, + self.image_std, + self.patch_size, + pool_w, + pool_h, + self.crop_mode, + ) + patch_batches.append(crops) + pooling_batches.append(pooled_idx) + grids.append(image_grid) + image_num_crops.append(int(crops.shape[0])) + + pixel_values = np.concatenate(patch_batches, axis=0) if patch_batches else np.zeros((0, 0, 0), dtype=np.float32) + image_token_pooling = ( + np.concatenate(pooling_batches, axis=0) if pooling_batches else np.zeros((0, pool_h * pool_w), dtype=np.int64) + ) + image_grids = np.concatenate(grids, axis=0) if grids else np.zeros((0, 4), dtype=np.int64) + image_num_crops_arr = np.asarray(image_num_crops, dtype=np.int64) + + return { + "pixel_values": pixel_values.astype(np.float32), + "image_token_pooling": image_token_pooling.astype(np.int64), + "image_grids": image_grids.astype(np.int64), + "image_num_crops": image_num_crops_arr, + } + + +__all__ = ["MolmoAct2ImageProcessor"] \ No newline at end of file From 1e6c40b1ec483c557766fd506a5b469757865651 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Sun, 12 Jul 2026 14:37:07 +0200 Subject: [PATCH 03/22] wip: molmo-preprocessors --- .../inference/preprocessors/molmoact2.py | 111 ++++++-- .../preprocessors/molmoact2_inputs.py | 248 ++++++++++++++++++ 2 files changed, 330 insertions(+), 29 deletions(-) create mode 100644 src/physicalai/inference/preprocessors/molmoact2_inputs.py diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index e10064ae..faa55955 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -14,9 +14,16 @@ import numpy as np -from physicalai.inference.constants import IMAGE_MASKS, IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.constants import IMAGES, STATE, TASK from .base import Preprocessor +from .molmoact2_image import MolmoAct2ImageProcessor as NumpyImagePatchifier +from .molmoact2_inputs import ( + MolmoAct2InputConfig, + build_batched_images, + default_action_dim_is_pad, + expand_image_placeholders, +) ACTION_OUTPUT_TOKEN = "" SETUP_START_TOKEN = "" @@ -121,6 +128,8 @@ def __init__( add_control_tokens: bool = False, state_stats: dict[str, list[float] | np.ndarray] | None = None, image_keys: list[str] | None = None, + image_processor_config: dict[str, Any] | None = None, + model_input_config: dict[str, Any] | None = None, processor_config: dict[str, Any] | None = None, ) -> None: self.tokenizer_name_or_path = tokenizer_name_or_path @@ -133,6 +142,20 @@ def __init__( self._tokenizer: Any = None del processor_config + image_cfg = dict(image_processor_config or {}) + self._image_patchifier = NumpyImagePatchifier( + size=image_cfg.get("size"), + image_mean=image_cfg.get("image_mean"), + image_std=image_cfg.get("image_std"), + crop_mode=str(image_cfg.get("crop_mode", "resize")), + patch_size=int(image_cfg.get("patch_size", 14)), + pooling_size=image_cfg.get("pooling_size"), + ) + + self._input_config = ( + MolmoAct2InputConfig(**dict(model_input_config)) if model_input_config is not None else None + ) + self._state_q01: np.ndarray | None = None self._state_q99: np.ndarray | None = None self._state_mask: np.ndarray | None = None @@ -300,28 +323,67 @@ def _extract_images_by_example(self, inputs: dict[str, Any], batch_size: int) -> return images_by_example @staticmethod - def _pack_images(images_by_example: list[list[np.ndarray]]) -> tuple[np.ndarray, np.ndarray]: - batch_size = len(images_by_example) - if batch_size == 0: - return np.empty((0, 0, 3, 0, 0), dtype=np.float32), np.empty((0, 0), dtype=bool) - - num_images = len(images_by_example[0]) - if any(len(example_images) != num_images for example_images in images_by_example): - msg = "MolmoAct2 requires a consistent number of images per batch element." + def _stack_flat_images(flat_images: list[np.ndarray]) -> np.ndarray: + """Stack example-major ``(C, H, W)`` crops into a ``(M, C, H, W)`` batch.""" + if not flat_images: + msg = "MolmoAct2 inference preprocessor requires at least one image input." raise ValueError(msg) + return np.stack([np.asarray(image, dtype=np.float32) for image in flat_images], axis=0) - if num_images == 0: - return np.empty((0, batch_size, 3, 0, 0), dtype=np.float32), np.empty((0, batch_size), dtype=bool) - - image_slots: list[np.ndarray] = [] - image_masks: list[np.ndarray] = [] - for image_idx in range(num_images): - slot_images = [images_by_example[batch_idx][image_idx].astype(np.float32, copy=False) for batch_idx in range(batch_size)] - image_slots.append(np.stack(slot_images, axis=0)) - image_masks.append(np.ones((batch_size,), dtype=bool)) - return np.stack(image_slots, axis=0), np.stack(image_masks, axis=0) + def _build_model_inputs( + self, + input_ids: np.ndarray, + attention_mask: np.ndarray, + flat_images: list[np.ndarray], + batch_size: int, + ) -> dict[str, np.ndarray]: + """Patchify images and assemble backbone-ready model inputs. + + Returns: + The exact tensor set the exported model consumes: ``input_ids``, + ``attention_mask``, ``token_type_ids``, ``images``, ``token_pooling`` + and ``action_dim_is_pad``. + """ + assert self._input_config is not None # noqa: S101 (validated by caller) + + image_out = self._image_patchifier(self._stack_flat_images(flat_images)) + pixel_values = np.asarray(image_out["pixel_values"], dtype=np.float32) + image_token_pooling = np.asarray(image_out["image_token_pooling"], dtype=np.int64) + image_grids = np.asarray(image_out["image_grids"], dtype=np.int64) + image_num_crops = np.asarray(image_out["image_num_crops"], dtype=np.int64) + + input_ids, attention_mask, token_type_ids = expand_image_placeholders( + config=self._input_config, + input_ids=input_ids, + attention_mask=attention_mask, + image_grids=image_grids, + ) + images, token_pooling = build_batched_images( + self._input_config, + input_ids, + pixel_values, + image_token_pooling, + image_grids, + image_num_crops, + ) + action_dim_is_pad = default_action_dim_is_pad(self._input_config, batch_size=batch_size) + + model_inputs: dict[str, np.ndarray] = { + "input_ids": input_ids.astype(np.int64), + "attention_mask": attention_mask.astype(np.int64), + "images": images.astype(np.float32), + "token_pooling": token_pooling.astype(np.int64), + "action_dim_is_pad": action_dim_is_pad, + } + if token_type_ids is not None: + model_inputs["token_type_ids"] = token_type_ids.astype(np.int64) + return model_inputs def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.ndarray]: + if self._input_config is None: + msg = "MolmoAct2Preprocessor requires model_input_config to build model inputs." + raise ValueError(msg) + inputs_dict = dict(inputs) state = self._extract_state(inputs_dict) @@ -346,9 +408,6 @@ def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.nd ) ) - del flat_images - images, image_masks = self._pack_images(images_by_example) - text_inputs = self.tokenizer(prompt_texts, padding=True) input_ids = np.asarray(text_inputs["input_ids"], dtype=np.int64) attention_mask = np.asarray(text_inputs["attention_mask"], dtype=np.int64) @@ -357,10 +416,4 @@ def __call__(self, inputs: dict[str, np.ndarray | list[str]]) -> dict[str, np.nd pad_token_id = self.tokenizer.pad_token_id input_ids, attention_mask = self._insert_bos(input_ids, attention_mask, int(bos_token_id), int(pad_token_id)) - return { - TOKENIZED_PROMPT: input_ids, - TOKENIZED_PROMPT_MASK: attention_mask, - STATE: state, - IMAGES: images, - IMAGE_MASKS: image_masks, - } + return self._build_model_inputs(input_ids, attention_mask, flat_images, batch_size) diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py new file mode 100644 index 00000000..62ace015 --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -0,0 +1,248 @@ +# 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, field + +import numpy as np + + +@dataclass +class MolmoAct2InputConfig: + """Token ids and layout flags needed to assemble MolmoAct2 model inputs.""" + + 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 = False + use_single_crop_start_token: bool = True + max_action_dim: int = 32 + env_action_dim: int = 0 + _image_token_ids: list[int] = field(default_factory=list, init=False, repr=False) + + def __post_init__(self) -> None: + 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, + ] + self._image_token_ids = [int(token_id) for token_id in ids if token_id is not None] + + @property + def image_token_ids(self) -> list[int]: + """Token ids that mark image content (for token type ids).""" + return self._image_token_ids + + +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.""" + 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 = 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.""" + 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.""" + if int(image_grids.shape[0]) == 0: + return input_ids, attention_mask, _build_token_type_ids(config, input_ids, attention_mask) + + pad_values = input_ids[attention_mask == 0] + pad_token_id = int(pad_values[0]) if pad_values.size > 0 else 0 + placeholder_id = int(config.image_placeholder_token_id) + + expanded_rows: list[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) + + max_len = max((len(row) for row in expanded_rows), default=1) + out_ids = np.full((len(expanded_rows), max_len), 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) + + +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)``. + """ + counts = (input_ids == int(config.image_end_token_id)).sum(1) # images per example + num_images = int(image_grids.shape[0]) + if int(counts.sum()) != num_images: + msg = f"image_end tokens ({int(counts.sum())}) do not match image grids ({num_images})." + raise ValueError(msg) + + num_examples = counts.shape[0] + n_crops, n_patches, pixels_per_patch = pixel_values.shape + del n_crops + + 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), counts) + 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) + patches_per_image = image_num_crops.astype(np.int64) * n_patches + + max_crops = int(crops_per_example.max()) if num_examples > 0 else 0 + images = np.full( + (num_examples, max_crops, n_patches, pixels_per_patch), + -1.0, + dtype=pixel_values.dtype, + ) + max_pooled = int(pooled_per_example.max()) if num_examples > 0 else 0 + token_pooling = np.full( + (num_examples, max_pooled, image_token_pooling.shape[-1]), + -1, + dtype=image_token_pooling.dtype, + ) + + crop_offset = 0 + pooled_offset = 0 + image_offset = 0 + for example_idx in range(num_examples): + num_example_images = int(counts[example_idx]) + num_example_crops = int(crops_per_example[example_idx]) + images[example_idx, :num_example_crops] = pixel_values[crop_offset : crop_offset + num_example_crops] + + example_pooling = image_token_pooling[ + pooled_offset : pooled_offset + int(pooled_per_example[example_idx]) + ].copy() + patch_offset = 0 + row = 0 + for local_image in range(num_example_images): + num_pooled = int(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(patches_per_image[image_offset + local_image]) + row += num_pooled + token_pooling[example_idx, : example_pooling.shape[0]] = example_pooling + + crop_offset += num_example_crops + pooled_offset += int(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.""" + 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 + + +__all__ = [ + "MolmoAct2InputConfig", + "build_batched_images", + "default_action_dim_is_pad", + "expand_image_placeholders", +] From 00a5512757fd1da912ddd9f3fa0119be3b8ebff3 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:39:47 +0200 Subject: [PATCH 04/22] wip: molmo pre and post processors --- src/physicalai/inference/component_factory.py | 3 + .../inference/postprocessors/__init__.py | 2 + .../inference/postprocessors/molmoact2.py | 65 +++ .../postprocessors/stats_denormalizer.py | 14 +- .../inference/preprocessors/__init__.py | 3 + .../inference/preprocessors/molmoact2.py | 439 ++++++++++++++++++ .../preprocessors/stats_normalizer.py | 14 +- .../postprocessors/test_molmoact2.py | 44 ++ .../postprocessors/test_stats_denormalizer.py | 17 + .../inference/preprocessors/test_molmoact2.py | 141 ++++++ .../preprocessors/test_stats_normalizer.py | 17 + 11 files changed, 747 insertions(+), 12 deletions(-) create mode 100644 src/physicalai/inference/postprocessors/molmoact2.py create mode 100644 src/physicalai/inference/preprocessors/molmoact2.py create mode 100644 tests/unit/inference/postprocessors/test_molmoact2.py create mode 100644 tests/unit/inference/preprocessors/test_molmoact2.py diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index e61f09dd..737cdade 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -105,6 +105,8 @@ 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("to_float_tensor", "physicalai.inference.preprocessors.ToFloatTensorPreprocessor") @@ -112,6 +114,7 @@ def __repr__(self) -> str: # 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") def resolve_artifact(spec: ComponentSpec, export_dir: Path) -> ComponentSpec: diff --git a/src/physicalai/inference/postprocessors/__init__.py b/src/physicalai/inference/postprocessors/__init__.py index 64a45d42..e4e52438 100644 --- a/src/physicalai/inference/postprocessors/__init__.py +++ b/src/physicalai/inference/postprocessors/__init__.py @@ -9,11 +9,13 @@ 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.molmoact2 import MolmoAct2Postprocessor from physicalai.inference.postprocessors.stats_denormalizer import StatsDenormalizer __all__ = [ "ActionChunkTrimmer", "ActionNormalizer", + "MolmoAct2Postprocessor", "Postprocessor", "StatsDenormalizer", ] diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py new file mode 100644 index 00000000..ce596d19 --- /dev/null +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -0,0 +1,65 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy postprocessor for MolmoAct2 exported models.""" + +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 + + +class MolmoAct2Postprocessor(Postprocessor): + """Clamp, denormalize, and optionally transform MolmoAct2 actions.""" + + def __init__( + self, + *, + action_stats: dict[str, Any] | None = None, + adapt_to_so101: bool = False, + joint_signs: list[float] | None = None, + joint_offsets: list[float] | None = None, + ) -> None: + signs = joint_signs or [] + offsets = joint_offsets or [] + if len(signs) != len(offsets): + msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" + raise ValueError(msg) + self._adapt_to_so101 = adapt_to_so101 + self._joint_signs = np.asarray(signs, dtype=np.float32) + self._joint_offsets = np.asarray(offsets, dtype=np.float32) + self._denormalizer = ( + StatsDenormalizer(stats={ACTION: action_stats}, mode="quantiles", features=[ACTION]) + if action_stats + else None + ) + + @override + def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + result = dict(outputs) + action = result.get(ACTION, result.get("actions")) + if action is None: + msg = "MolmoAct2 postprocessor expected an action tensor" + raise ValueError(msg) + action = np.clip(np.asarray(action), -1.0, 1.0) + if self._denormalizer is not None: + action = self._denormalizer({ACTION: action})[ACTION] + if self._adapt_to_so101: + count = min(self._joint_signs.size, action.shape[-1]) + transformed = np.array(action, copy=True) + transformed[..., :count] = self._joint_signs[:count] * ( + action[..., :count] - self._joint_offsets[:count] + ) + action = transformed + result.pop("actions", None) + 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 e19e07fd..173f83b8 100644 --- a/src/physicalai/inference/postprocessors/stats_denormalizer.py +++ b/src/physicalai/inference/postprocessors/stats_denormalizer.py @@ -185,21 +185,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 35cfe2e4..2347439f 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -10,6 +10,7 @@ from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer 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 @@ -21,6 +22,8 @@ __all__ = [ "HFTokenizer", "LambdaPreprocessor", + "MolmoAct2ModelInputs", + "MolmoAct2Preprocessor", "NewLinePreprocessor", "OVTokenizer", "Pi05Preprocessor", diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py new file mode 100644 index 00000000..e9d0035d --- /dev/null +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -0,0 +1,439 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy preprocessors for MolmoAct2 exported models.""" + +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, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.preprocessors.base import Preprocessor +from physicalai.inference.preprocessors.stats_normalizer import StatsNormalizer + +_STATE_START_TOKEN = "" # noqa: S105 +_STATE_END_TOKEN = "" # noqa: S105 +_STATE_TOKEN_PREFIX = " np.ndarray: + """Apply the MolmoAct2 joint-frame transform to leading dimensions.""" + count = min(signs.size, values.shape[-1]) + output = np.array(values, copy=True) + joints = values[..., :count] + output[..., :count] = signs[:count] * (joints - offsets[:count]) if inverse else signs[:count] * joints + offsets[:count] + return output + + +def _normalize_text(text: str) -> str: + normalized = re.sub(r"\s+", " ", str(text or "")).strip() + if not normalized: + return "" + for pattern in _PREFIX_PATTERNS: + normalized = pattern.sub("", normalized, count=1).strip() + return normalized.rstrip(_TRAILING_PUNCTUATION).strip().lower() + + +def _discrete_state_string(state: np.ndarray, num_state_tokens: int) -> str: + values = np.nan_to_num(np.asarray(state, dtype=np.float32), nan=0.0, posinf=1.0, neginf=-1.0) + values = np.clip(values, -1.0, 1.0) + token_ids = np.rint((values + 1.0) / 2.0 * (num_state_tokens - 1)).astype(np.int64) + payload = "".join(f"{_STATE_TOKEN_PREFIX}{int(token_id)}>" for token_id in token_ids.reshape(-1)) + return f"{_STATE_START_TOKEN}{payload}{_STATE_END_TOKEN}" + + +def _wrapped_text(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 _robot_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 = _wrapped_text(setup_type, "", "", enabled=add_setup_tokens) + control = _wrapped_text(control_mode, "", "", enabled=add_control_tokens) + discrete_state = _discrete_state_string(state, num_state_tokens) + prompt = ( + f"The task is to {task}. The setup is {setup}. " + f"The current state of the robot is {discrete_state}. " + 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|>" + elif num_images > 1: + image_prefix = "".join(f"Image {index + 1}<|image|>" for index in range(num_images)) + else: + image_prefix = "" + return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{_ACTION_OUTPUT_TOKEN}" + + +class MolmoAct2Preprocessor(Preprocessor): + """Prepare MolmoAct2 prompts and images before tokenization.""" + + def __init__( + self, + *, + image_keys: list[str], + state_stats: dict[str, Any] | None = None, + 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, + adapt_to_so101: bool = False, + joint_signs: list[float] | None = None, + joint_offsets: list[float] | None = None, + ) -> None: + if num_state_tokens <= 0: + msg = f"num_state_tokens must be > 0, got {num_state_tokens}" + raise ValueError(msg) + signs = joint_signs or [] + offsets = joint_offsets or [] + if len(signs) != len(offsets): + msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" + 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._adapt_to_so101 = adapt_to_so101 + self._joint_signs = np.asarray(signs, dtype=np.float32) + self._joint_offsets = np.asarray(offsets, dtype=np.float32) + self._normalizer = ( + StatsNormalizer(stats={STATE: state_stats}, mode="quantiles", features=[STATE]) if state_stats else None + ) + + @override + def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: + 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._adapt_to_so101: + state = _joint_transform(state, self._joint_signs, self._joint_offsets, inverse=False) + if self._normalizer is not None: + state = self._normalizer({STATE: state})[STATE] + state = np.clip(state, -1.0, 1.0) + + images = self._extract_images(outputs, batch_size=state.shape[0]) + tasks = self._extract_tasks(outputs, batch_size=state.shape[0]) + outputs[IMAGES] = np.stack([self._resize_image(image) for image in images], axis=0) + outputs[TASK] = [ + _robot_prompt( + task=tasks[index], + state=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 _extract_images(self, inputs: dict[str, Any], *, batch_size: int) -> list[np.ndarray]: + images_value = inputs.get(IMAGES) + images: list[np.ndarray] = [] + if self._image_keys: + for name in self._image_keys: + flat_key = name if name.startswith(f"{IMAGES}.") else f"{IMAGES}.{name}" + if flat_key in inputs: + images.append(np.asarray(inputs[flat_key])) + elif isinstance(images_value, dict) and name.removeprefix(f"{IMAGES}.") in images_value: + images.append(np.asarray(images_value[name.removeprefix(f"{IMAGES}.")])) + elif isinstance(images_value, np.ndarray): + images = [images_value] + elif isinstance(images_value, dict): + images = [np.asarray(value) for key, value in images_value.items() if "is_pad" not in str(key)] + else: + 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) + for image in images: + if image.ndim != _IMAGE_NDIM or image.shape[1] != _NUM_CHANNELS: + msg = f"Expected BCHW image with 3 channels, got {image.shape}" + raise ValueError(msg) + if image.shape[0] != batch_size: + msg = f"Image batch size mismatch: expected {batch_size}, got {image.shape[0]}" + raise ValueError(msg) + return images + + @staticmethod + def _extract_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) + if isinstance(source, str): + tasks = [source] * batch_size + else: + tasks = [str(value) for value in 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(task) for task in tasks] + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + height, width = self._image_size + output: list[np.ndarray] = [] + for sample in image: + if sample.dtype == np.uint8: + pixels = sample + elif np.issubdtype(sample.dtype, np.floating): + float_pixels = sample.astype(np.float32) + if float(np.max(float_pixels)) <= 1.0: + float_pixels *= 255.0 + pixels = np.clip(float_pixels, 0.0, 255.0).astype(np.uint8) + else: + msg = f"Unsupported image dtype: {sample.dtype}" + raise ValueError(msg) + hwc = np.transpose(pixels, (1, 2, 0)) + resized = cv2.resize(hwc, (width, height), interpolation=cv2.INTER_LINEAR_EXACT) + output.append(np.transpose(resized, (2, 0, 1)).astype(np.float32) / 255.0) + return np.stack(output, axis=0) + + +class MolmoAct2ModelInputs(Preprocessor): + """Assemble tokenized prompts and packed images into MolmoAct2 model inputs.""" + + def __init__( + self, + *, + 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, + 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_use_col_tokens: bool = True, + use_single_crop_col_tokens: bool = False, + use_single_crop_start_token: bool = True, + image_token_ids: list[int] | None = None, + ) -> None: + self._max_action_dim = max_action_dim + self._action_dim = action_dim + self._bos_token_id = bos_token_id + self._pad_token_id = pad_token_id + self._placeholder_id = image_placeholder_token_id + self._image_start_id = image_start_token_id + self._image_end_id = image_end_token_id + self._image_patch_id = image_patch_id + self._image_col_id = image_col_id + self._low_res_start_id = low_res_image_start_token_id or image_start_token_id + self._height, self._width = image_size + self._patch_size = patch_size + self._pool_h, self._pool_w = pooling_size + self._mean = np.asarray(image_mean or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) + self._std = np.asarray(image_std or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) + self._image_use_col_tokens = image_use_col_tokens + self._use_single_crop_col_tokens = use_single_crop_col_tokens + self._use_single_crop_start_token = use_single_crop_start_token + self._image_token_ids = np.asarray(image_token_ids or [], dtype=np.int64) + self._pooling, self._pooled_h, self._pooled_w = self._pooling_indices() + + @override + def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: + 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 + if channels != _NUM_CHANNELS or (height, width) != (self._height, self._width): + msg = f"Unexpected packed image shape {images.shape}" + raise ValueError(msg) + + flat_images = images.transpose(1, 0, 2, 3, 4).reshape(batch_size * num_images, channels, height, width) + pixel_values = self._patchify((flat_images - self._mean) / self._std) + grids = np.tile(np.array([[self._pooled_h, self._pooled_w, 0, 0]], dtype=np.int64), (batch_size * num_images, 1)) + input_ids, attention_mask = self._expand_placeholders(input_ids, attention_mask, grids) + token_type_ids = self._token_type_ids(input_ids, attention_mask) + batched_images = pixel_values.reshape(batch_size, num_images, pixel_values.shape[1], pixel_values.shape[2]) + + pooling = [] + patches_per_image = pixel_values.shape[1] + for image_index in range(num_images): + block = np.where(self._pooling >= 0, self._pooling + image_index * patches_per_image, self._pooling) + pooling.append(block) + token_pooling = np.tile(np.concatenate(pooling, axis=0)[None, ...], (batch_size, 1, 1)) + action_dim_is_pad = np.ones((batch_size, self._max_action_dim), dtype=np.bool_) + action_dim_is_pad[:, : self._action_dim] = False + + outputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + IMAGES: batched_images.astype(np.float32), + "token_pooling": token_pooling.astype(np.int64), + "action_dim_is_pad": action_dim_is_pad, + } + 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: list[np.ndarray] = [] + for row_ids, row_mask in zip(ids, mask, strict=True): + valid_ids = row_ids[row_mask.astype(np.bool_)] + if valid_ids.size == 0 or valid_ids[0] != self._bos_token_id: + valid_ids = np.concatenate((np.array([self._bos_token_id], dtype=ids.dtype), valid_ids)) + rows.append(valid_ids) + width = max((row.size for row in rows), default=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 + + 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, + ) + pooling = indices.reshape(pooled_h, self._pool_h, pooled_w, self._pool_w) + return pooling.transpose(0, 2, 1, 3).reshape(-1, self._pool_h * self._pool_w), pooled_h, pooled_w + + def _image_sequence(self, grid: np.ndarray) -> list[int]: + resized_h, resized_w, height, width = (int(value) for value in grid) + + def rows(row_count: int, col_count: int, *, use_col: bool) -> list[int]: + row = [self._image_patch_id] * col_count + if use_col and self._image_col_id is not None: + row.append(self._image_col_id) + return row * row_count + + if height == 0 or width == 0: + return [ + self._image_start_id, + *rows(resized_h, resized_w, use_col=self._use_single_crop_col_tokens), + self._image_end_id, + ] + low_start = self._low_res_start_id if self._use_single_crop_start_token else self._image_start_id + return [ + low_start, + *rows(resized_h, resized_w, use_col=self._use_single_crop_col_tokens), + self._image_end_id, + self._image_start_id, + *rows(height, width, use_col=self._image_use_col_tokens), + self._image_end_id, + ] + + def _expand_placeholders( + self, + ids: np.ndarray, + mask: np.ndarray, + grids: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + rows: list[np.ndarray] = [] + grid_index = 0 + for row_ids, row_mask in zip(ids, mask, strict=True): + expanded: list[int] = [] + for token in row_ids[row_mask.astype(np.bool_)]: + if int(token) == self._placeholder_id: + if grid_index >= grids.shape[0]: + msg = "Not enough image grids to expand all <|image|> placeholders" + raise ValueError(msg) + expanded.extend(self._image_sequence(grids[grid_index])) + grid_index += 1 + else: + expanded.append(int(token)) + rows.append(np.asarray(expanded, dtype=ids.dtype)) + if grid_index != grids.shape[0]: + msg = f"Image placeholders ({grid_index}) do not match images ({grids.shape[0]})" + raise ValueError(msg) + width = max((row.size for row in rows), default=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 + + def _token_type_ids(self, ids: np.ndarray, mask: np.ndarray) -> np.ndarray | None: + if self._image_token_ids.size == 0: + return None + return (np.isin(ids, self._image_token_ids) & mask.astype(np.bool_)).astype(np.int64) + + +__all__ = ["MolmoAct2ModelInputs", "MolmoAct2Preprocessor"] diff --git a/src/physicalai/inference/preprocessors/stats_normalizer.py b/src/physicalai/inference/preprocessors/stats_normalizer.py index 2d4ba914..dbd2bb70 100644 --- a/src/physicalai/inference/preprocessors/stats_normalizer.py +++ b/src/physicalai/inference/preprocessors/stats_normalizer.py @@ -186,24 +186,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)): @@ -211,6 +212,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_molmoact2.py b/tests/unit/inference/postprocessors/test_molmoact2.py new file mode 100644 index 00000000..972f223d --- /dev/null +++ b/tests/unit/inference/postprocessors/test_molmoact2.py @@ -0,0 +1,44 @@ +# 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 MolmoAct2Postprocessor + + +class TestMolmoAct2Postprocessor: + def test_clamps_masked_denormalizes_and_transforms(self) -> None: + processor = MolmoAct2Postprocessor( + action_stats={ + "q01": [0.0, 0.0, 0.0], + "q99": [2.0, 2.0, 2.0], + "mask": [True, False, True], + }, + adapt_to_so101=True, + joint_signs=[1.0, -1.0], + joint_offsets=[0.0, 2.0], + ) + + result = 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() + 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="action tensor"): + MolmoAct2Postprocessor()({"other": np.zeros(1)}) + + def test_registry_alias_instantiates(self) -> None: + processor = instantiate_component(ComponentSpec(type="molmoact2_postprocess")) + assert isinstance(processor, MolmoAct2Postprocessor) 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_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py new file mode 100644 index 00000000..7156f4a6 --- /dev/null +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -0,0 +1,141 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest + +from physicalai.inference.constants import IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.component_factory import instantiate_component +from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor + + +def _raw_preprocessor(**kwargs) -> MolmoAct2Preprocessor: + return MolmoAct2Preprocessor( + image_keys=["top", "wrist"], + image_size=(28, 28), + num_state_tokens=4, + setup_type="tabletop", + control_mode="joint", + **kwargs, + ) + + +def _model_inputs(**kwargs) -> MolmoAct2ModelInputs: + return MolmoAct2ModelInputs( + 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_size=(28, 28), + patch_size=14, + pooling_size=(2, 2), + image_token_ids=[10, 11, 12, 13], + **kwargs, + ) + + +class TestMolmoAct2Preprocessor: + def test_builds_prompt_and_packs_ordered_cameras(self) -> None: + processor = _raw_preprocessor() + inputs = { + STATE: np.array([[-1.0, 1.0]], dtype=np.float32), + TASK: ["Task: Pick up."], + f"{IMAGES}.wrist": np.full((1, 3, 28, 28), 255, dtype=np.uint8), + f"{IMAGES}.top": np.zeros((1, 3, 28, 28), dtype=np.uint8), + } + + result = processor(inputs) + + 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_applies_masked_normalization_and_joint_transform(self) -> None: + processor = MolmoAct2Preprocessor( + image_keys=[], + image_size=(28, 28), + state_stats={"q01": [0.0, 0.0], "q99": [2.0, 2.0], "mask": [True, False]}, + adapt_to_so101=True, + joint_signs=[1.0, -1.0], + joint_offsets=[0.0, 2.0], + ) + result = processor( + { + STATE: np.array([[1.0, 1.0]], dtype=np.float32), + TASK: "move", + IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8), + }, + ) + assert "" in result[TASK][0] + + def test_rejects_missing_state(self) -> None: + with pytest.raises(ValueError, match="state"): + _raw_preprocessor()({TASK: ["move"], IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8)}) + + def test_registry_alias_instantiates(self) -> None: + processor = instantiate_component(ComponentSpec(type="molmoact2", image_keys=[])) + assert isinstance(processor, MolmoAct2Preprocessor) + + +class TestMolmoAct2ModelInputs: + def test_assembles_model_inputs(self) -> None: + processor = _model_inputs() + result = processor( + { + TOKENIZED_PROMPT: np.array([[99, 5, 99, 0]], dtype=np.int64), + TOKENIZED_PROMPT_MASK: np.array([[1, 1, 1, 0]], dtype=np.bool_), + IMAGES: np.zeros((2, 1, 3, 28, 28), dtype=np.float32), + }, + ) + + assert set(result) == { + "input_ids", + "attention_mask", + "token_type_ids", + "images", + "token_pooling", + "action_dim_is_pad", + } + assert int(result["input_ids"][0, 0]) == 1 + assert result["images"].shape == (1, 2, 4, 588) + assert result["token_pooling"].shape == (1, 2, 4) + np.testing.assert_array_equal(result["action_dim_is_pad"], [[False, False, True, True]]) + assert result["token_type_ids"].sum() > 0 + + def test_rejects_placeholder_image_mismatch(self) -> None: + with pytest.raises(ValueError, match="placeholders"): + _model_inputs()( + { + TOKENIZED_PROMPT: np.array([[99, 5]], dtype=np.int64), + TOKENIZED_PROMPT_MASK: np.ones((1, 2), dtype=np.bool_), + IMAGES: np.zeros((2, 1, 3, 28, 28), dtype=np.float32), + }, + ) + + def test_registry_alias_instantiates(self) -> None: + spec = ComponentSpec( + type="molmoact2_inputs", + 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, + ) + assert isinstance(instantiate_component(spec), MolmoAct2ModelInputs) diff --git a/tests/unit/inference/preprocessors/test_stats_normalizer.py b/tests/unit/inference/preprocessors/test_stats_normalizer.py index 440b21b8..7986489a 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": [-1.0, -1.0, -1.0], + "q99": [1.0, 1.0, 1.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, -0.5]]) + class TestStatsNormalizerIdentity: def test_identity_mode_passthrough(self, stats_dir: Path) -> None: From a9819203f7e5fd04130ca735f96882361d983a2f Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:57:54 +0100 Subject: [PATCH 05/22] wip: working pre / post processors --- .../inference/preprocessors/molmoact2.py | 9 +- .../inference/preprocessors/test_molmoact2.py | 86 +++++++++++++++++++ uv.lock | 2 +- 3 files changed, 93 insertions(+), 4 deletions(-) diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index e9d0035d..4964e4a7 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -340,7 +340,7 @@ def _insert_bos(self, ids: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np if valid_ids.size == 0 or valid_ids[0] != self._bos_token_id: valid_ids = np.concatenate((np.array([self._bos_token_id], dtype=ids.dtype), valid_ids)) rows.append(valid_ids) - width = max((row.size for row in rows), default=1) + 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): @@ -406,10 +406,12 @@ def _expand_placeholders( grids: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: rows: list[np.ndarray] = [] + expanded_widths: list[int] = [] grid_index = 0 for row_ids, row_mask in zip(ids, mask, strict=True): + valid = row_mask.astype(np.bool_) expanded: list[int] = [] - for token in row_ids[row_mask.astype(np.bool_)]: + for token in row_ids[valid]: if int(token) == self._placeholder_id: if grid_index >= grids.shape[0]: msg = "Not enough image grids to expand all <|image|> placeholders" @@ -419,10 +421,11 @@ def _expand_placeholders( else: expanded.append(int(token)) rows.append(np.asarray(expanded, dtype=ids.dtype)) + expanded_widths.append(len(expanded) + int((~valid).sum())) if grid_index != grids.shape[0]: msg = f"Image placeholders ({grid_index}) do not match images ({grids.shape[0]})" raise ValueError(msg) - width = max((row.size for row in rows), default=1) + width = max(expanded_widths, default=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): diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 7156f4a6..0ed33c80 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -114,6 +114,19 @@ def test_assembles_model_inputs(self) -> None: np.testing.assert_array_equal(result["action_dim_is_pad"], [[False, False, True, True]]) assert result["token_type_ids"].sum() > 0 + def test_preserves_masked_tokenizer_padding(self) -> None: + result = _model_inputs()( + { + 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_rejects_placeholder_image_mismatch(self) -> None: with pytest.raises(ValueError, match="placeholders"): _model_inputs()( @@ -139,3 +152,76 @@ def test_registry_alias_instantiates(self) -> None: low_res_image_start_token_id=10, ) assert isinstance(instantiate_component(spec), MolmoAct2ModelInputs) + + +class TestMolmoAct2ManifestPipeline: + def test_processes_observation_and_action(self, monkeypatch) -> None: + from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer + + class StubTokenizer: + 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: StubTokenizer(), + ) + 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, + 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), + } + for spec in specs: + values = instantiate_component(spec)(values) + + assert isinstance(instantiate_component(specs[1]), HFTokenizer) + 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_stats={"q01": [0.0, 0.0], "q99": [2.0, 2.0]}, + ), + ) + 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/uv.lock b/uv.lock index f9ba9ee9..d197015d 100644 --- a/uv.lock +++ b/uv.lock @@ -1777,7 +1777,7 @@ requires-dist = [ { name = "msgpack", marker = "extra == 'transport'", specifier = "==1.2.1" }, { name = "mypy", marker = "extra == 'tests'" }, { name = "num2words", specifier = ">=0.5.14,<0.6.0" }, - { name = "numpy", specifier = ">=1.24" }, + { name = "numpy" }, { name = "onnxruntime" }, { name = "opencv-python-headless" }, { name = "opencv-python-headless", marker = "extra == 'basler'" }, From 0b8e35dff21f5f9e85b037ec3ac72048faeeb282 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:58:27 +0100 Subject: [PATCH 06/22] wip: working pre / post processors --- .../inference/postprocessors/molmoact2.py | 4 +--- .../inference/preprocessors/molmoact2.py | 8 ++++++-- .../preprocessors/molmoact2_image.py | 12 +++++------ .../preprocessors/molmoact2_inputs.py | 20 +++++++++---------- 4 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py index ce596d19..2def501d 100644 --- a/src/physicalai/inference/postprocessors/molmoact2.py +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -53,9 +53,7 @@ def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: if self._adapt_to_so101: count = min(self._joint_signs.size, action.shape[-1]) transformed = np.array(action, copy=True) - transformed[..., :count] = self._joint_signs[:count] * ( - action[..., :count] - self._joint_offsets[:count] - ) + transformed[..., :count] = self._joint_signs[:count] * (action[..., :count] - self._joint_offsets[:count]) action = transformed result.pop("actions", None) result[ACTION] = action diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 4964e4a7..084d4bc4 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -44,7 +44,9 @@ def _joint_transform( count = min(signs.size, values.shape[-1]) output = np.array(values, copy=True) joints = values[..., :count] - output[..., :count] = signs[:count] * (joints - offsets[:count]) if inverse else signs[:count] * joints + offsets[:count] + output[..., :count] = ( + signs[:count] * (joints - offsets[:count]) if inverse else signs[:count] * joints + offsets[:count] + ) return output @@ -308,7 +310,9 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: flat_images = images.transpose(1, 0, 2, 3, 4).reshape(batch_size * num_images, channels, height, width) pixel_values = self._patchify((flat_images - self._mean) / self._std) - grids = np.tile(np.array([[self._pooled_h, self._pooled_w, 0, 0]], dtype=np.int64), (batch_size * num_images, 1)) + grids = np.tile( + np.array([[self._pooled_h, self._pooled_w, 0, 0]], dtype=np.int64), (batch_size * num_images, 1) + ) input_ids, attention_mask = self._expand_placeholders(input_ids, attention_mask, grids) token_type_ids = self._token_type_ids(input_ids, attention_mask) batched_images = pixel_values.reshape(batch_size, num_images, pixel_values.shape[1], pixel_values.shape[2]) diff --git a/src/physicalai/inference/preprocessors/molmoact2_image.py b/src/physicalai/inference/preprocessors/molmoact2_image.py index b110313e..bf340c9f 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_image.py +++ b/src/physicalai/inference/preprocessors/molmoact2_image.py @@ -34,9 +34,7 @@ def _resize_image(image: np.ndarray, desired_output_size: list[int]) -> np.ndarr def _select_tiling(h: int, w: int, patch_size: int, max_num_crops: int) -> np.ndarray: tilings: list[tuple[int, int]] = [] for i in range(1, max_num_crops + 1): - for j in range(1, max_num_crops + 1): - if i * j <= max_num_crops: - tilings.append((i, j)) + tilings.extend((i, j) for j in range(1, max_num_crops + 1) if i * j <= max_num_crops) tilings.sort(key=lambda x: (x[0] * x[1], x[0])) candidate_tilings = np.asarray(tilings, dtype=np.int32) candidate_resolutions = candidate_tilings * patch_size @@ -231,7 +229,7 @@ def _to_hwc_uint8(images_bchw: np.ndarray) -> list[np.ndarray]: img = image if np.issubdtype(img.dtype, np.floating): if float(np.max(img)) <= 1.0: - img = img * 255.0 + img *= 255.0 img = np.clip(img, 0.0, 255.0).astype(np.uint8) elif img.dtype != np.uint8: img = np.clip(img, 0, 255).astype(np.uint8) @@ -294,7 +292,9 @@ def __call__(self, images_bchw: np.ndarray) -> dict[str, np.ndarray]: pixel_values = np.concatenate(patch_batches, axis=0) if patch_batches else np.zeros((0, 0, 0), dtype=np.float32) image_token_pooling = ( - np.concatenate(pooling_batches, axis=0) if pooling_batches else np.zeros((0, pool_h * pool_w), dtype=np.int64) + np.concatenate(pooling_batches, axis=0) + if pooling_batches + else np.zeros((0, pool_h * pool_w), dtype=np.int64) ) image_grids = np.concatenate(grids, axis=0) if grids else np.zeros((0, 4), dtype=np.int64) image_num_crops_arr = np.asarray(image_num_crops, dtype=np.int64) @@ -307,4 +307,4 @@ def __call__(self, images_bchw: np.ndarray) -> dict[str, np.ndarray]: } -__all__ = ["MolmoAct2ImageProcessor"] \ No newline at end of file +__all__ = ["MolmoAct2ImageProcessor"] diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py index 62ace015..c0f6bc5f 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -72,28 +72,26 @@ def _image_token_ids_for_grid(config: MolmoAct2InputConfig, grid: np.ndarray) -> 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) + 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 = row + [image_col_id] + 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] - ) + 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] + 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] + 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 From 1789fb9acd0aa4f6184f6443382b27b834cd2081 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:31:20 +0100 Subject: [PATCH 07/22] wip: working asset tokenizer --- src/physicalai/inference/component_factory.py | 1 + .../inference/preprocessors/__init__.py | 2 + .../preprocessors/asset_tokenizer.py | 95 +++++++++++++++++++ .../preprocessors/test_asset_tokenizer.py | 91 ++++++++++++++++++ 4 files changed, 189 insertions(+) create mode 100644 src/physicalai/inference/preprocessors/asset_tokenizer.py create mode 100644 tests/unit/inference/preprocessors/test_asset_tokenizer.py diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index 47f23c8d..cfee1f12 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -105,6 +105,7 @@ 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("asset_tokenizer", "physicalai.inference.preprocessors.AssetTokenizer") 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") diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index 2347439f..ddeee9ed 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -7,6 +7,7 @@ flattens and filters them for the runtime adapter. """ +from physicalai.inference.preprocessors.asset_tokenizer import AssetTokenizer from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer from physicalai.inference.preprocessors.lambda_processor import LambdaPreprocessor @@ -20,6 +21,7 @@ from physicalai.inference.preprocessors.to_tensor import ToFloatTensorPreprocessor __all__ = [ + "AssetTokenizer", "HFTokenizer", "LambdaPreprocessor", "MolmoAct2ModelInputs", diff --git a/src/physicalai/inference/preprocessors/asset_tokenizer.py b/src/physicalai/inference/preprocessors/asset_tokenizer.py new file mode 100644 index 00000000..e712fad0 --- /dev/null +++ b/src/physicalai/inference/preprocessors/asset_tokenizer.py @@ -0,0 +1,95 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Tokenizer preprocessor loaded from a bundled tokenizer artifact.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +from physicalai.inference.constants import TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.preprocessors.base import Preprocessor + +_SUPPORTED_TOKENIZER_CLASSES = {"Qwen2Tokenizer"} + + +class AssetTokenizer(Preprocessor): + """Load an allowlisted Transformers tokenizer from a local artifact. + + Args: + artifact: Path to a bundled tokenizer file such as ``tokenizer.json``. + tokenizer_class: Allowlisted Transformers tokenizer class name. + tokenizer_options: Checkpoint-derived tokenizer construction options. + max_token_len: Maximum encoded prompt length. + """ + + def __init__( + self, + artifact: str, + tokenizer_class: str, + tokenizer_options: dict[str, Any] | None = None, + max_token_len: int = 512, + ) -> None: + """Initialize a tokenizer from a bundled local artifact. + + Raises: + FileNotFoundError: If the tokenizer artifact does not exist. + ImportError: If Transformers is not installed. + ValueError: If the requested tokenizer class is not supported. + """ + super().__init__() + artifact_path = Path(artifact) + if not artifact_path.is_file(): + msg = f"Tokenizer artifact does not exist: {artifact_path}" + raise FileNotFoundError(msg) + if tokenizer_class not in _SUPPORTED_TOKENIZER_CLASSES: + msg = f"Unsupported asset tokenizer class: {tokenizer_class!r}" + raise ValueError(msg) + + try: + import transformers # ruff: ignore[PLC0415] + except ImportError as exc: + msg = "Tokenizer requires transformers. Install with: pip install transformers" + raise ImportError(msg) from exc + + tokenizer_type = getattr(transformers, tokenizer_class) + self._tokenizer = tokenizer_type.from_pretrained( + artifact_path.parent, + local_files_only=True, + **(tokenizer_options or {}), + ) + self._max_token_len = max_token_len + + def __call__(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + """Tokenize tasks and replace them with token IDs and masks. + + Returns: + Input values with tasks replaced by token IDs and attention masks. + + Raises: + TypeError: If the task value is not a list. + """ + batch_tasks = inputs[TASK] + if not isinstance(batch_tasks, list): + msg = f"Expected TASK to be a list of strings, got {type(batch_tasks)}" + raise TypeError(msg) + + outputs = dict(inputs) + outputs.pop(TASK) + encoded_tokens = self._tokenizer( + batch_tasks, + max_length=self._max_token_len, + truncation=True, + padding="max_length", + return_tensors="np", + ) + outputs[TOKENIZED_PROMPT] = encoded_tokens["input_ids"] + outputs[TOKENIZED_PROMPT_MASK] = encoded_tokens["attention_mask"].astype(np.bool_) + return outputs + + def __repr__(self) -> str: + """Return string representation of the preprocessor.""" + return f"{self.__class__.__name__}(tokenizer={self._tokenizer.name_or_path!r})" diff --git a/tests/unit/inference/preprocessors/test_asset_tokenizer.py b/tests/unit/inference/preprocessors/test_asset_tokenizer.py new file mode 100644 index 00000000..c5f2353b --- /dev/null +++ b/tests/unit/inference/preprocessors/test_asset_tokenizer.py @@ -0,0 +1,91 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from physicalai.inference.component_factory import instantiate_component, resolve_artifact +from physicalai.inference.constants import TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.preprocessors import AssetTokenizer + + +def _mock_tokenizer() -> MagicMock: + tokenizer = MagicMock() + tokenizer.name_or_path = "local-tokenizer" + + def _encode(tasks, **kwargs): + length = kwargs["max_length"] + return { + "input_ids": np.ones((len(tasks), length), dtype=np.int64), + "attention_mask": np.ones((len(tasks), length), dtype=np.int64), + } + + tokenizer.side_effect = _encode + return tokenizer + + +def test_loads_allowlisted_tokenizer_with_dynamic_options(tmp_path: Path) -> None: + tokenizer_file = tmp_path / "tokenizer.json" + tokenizer_file.write_text("{}", encoding="utf-8") + transformers = MagicMock() + tokenizer = _mock_tokenizer() + transformers.Qwen2Tokenizer.from_pretrained.return_value = tokenizer + options = { + "bos_token": "<|im_end|>", + "extra_special_tokens": ["", "<|image|>"], + "model_max_length": 1010000, + } + + with patch.dict("sys.modules", {"transformers": transformers}): + preprocessor = AssetTokenizer( + artifact=str(tokenizer_file), + tokenizer_class="Qwen2Tokenizer", + tokenizer_options=options, + max_token_len=4, + ) + + transformers.Qwen2Tokenizer.from_pretrained.assert_called_once_with( + tmp_path, + local_files_only=True, + **options, + ) + result = preprocessor({TASK: ["pick up the block"]}) + assert result[TOKENIZED_PROMPT].shape == (1, 4) + assert result[TOKENIZED_PROMPT_MASK].dtype == np.bool_ + assert TASK not in result + + +def test_rejects_unsupported_tokenizer_class(tmp_path: Path) -> None: + tokenizer_file = tmp_path / "tokenizer.json" + tokenizer_file.write_text("{}", encoding="utf-8") + + with pytest.raises(ValueError, match="Unsupported asset tokenizer class"): + AssetTokenizer(artifact=str(tokenizer_file), tokenizer_class="ArbitraryTokenizer") + + +def test_manifest_resolves_flat_artifact(tmp_path: Path) -> None: + tokenizer_file = tmp_path / "tokenizer.json" + tokenizer_file.write_text("{}", encoding="utf-8") + transformers = MagicMock() + transformers.Qwen2Tokenizer.from_pretrained.return_value = _mock_tokenizer() + spec = resolve_artifact( + ComponentSpec( + type="asset_tokenizer", + artifact="tokenizer.json", + tokenizer_class="Qwen2Tokenizer", + tokenizer_options={"extra_special_tokens": ["<|image|>"]}, + ), + tmp_path, + ) + + with patch.dict("sys.modules", {"transformers": transformers}): + component = instantiate_component(spec) + + assert isinstance(component, AssetTokenizer) + assert spec.flat_params["artifact"] == str(tokenizer_file) \ No newline at end of file From 67491063233db4fe121d10c2b0a8c27300b750c1 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:35:52 +0100 Subject: [PATCH 08/22] wip: working ov tokenizer --- src/physicalai/inference/component_factory.py | 1 - .../inference/preprocessors/__init__.py | 2 - .../preprocessors/asset_tokenizer.py | 95 ------------------- .../preprocessors/test_asset_tokenizer.py | 91 ------------------ 4 files changed, 189 deletions(-) delete mode 100644 src/physicalai/inference/preprocessors/asset_tokenizer.py delete mode 100644 tests/unit/inference/preprocessors/test_asset_tokenizer.py diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index cfee1f12..47f23c8d 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -105,7 +105,6 @@ 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("asset_tokenizer", "physicalai.inference.preprocessors.AssetTokenizer") 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") diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index ddeee9ed..2347439f 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -7,7 +7,6 @@ flattens and filters them for the runtime adapter. """ -from physicalai.inference.preprocessors.asset_tokenizer import AssetTokenizer from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer from physicalai.inference.preprocessors.lambda_processor import LambdaPreprocessor @@ -21,7 +20,6 @@ from physicalai.inference.preprocessors.to_tensor import ToFloatTensorPreprocessor __all__ = [ - "AssetTokenizer", "HFTokenizer", "LambdaPreprocessor", "MolmoAct2ModelInputs", diff --git a/src/physicalai/inference/preprocessors/asset_tokenizer.py b/src/physicalai/inference/preprocessors/asset_tokenizer.py deleted file mode 100644 index e712fad0..00000000 --- a/src/physicalai/inference/preprocessors/asset_tokenizer.py +++ /dev/null @@ -1,95 +0,0 @@ -# Copyright (C) 2026 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -"""Tokenizer preprocessor loaded from a bundled tokenizer artifact.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import numpy as np - -from physicalai.inference.constants import TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK -from physicalai.inference.preprocessors.base import Preprocessor - -_SUPPORTED_TOKENIZER_CLASSES = {"Qwen2Tokenizer"} - - -class AssetTokenizer(Preprocessor): - """Load an allowlisted Transformers tokenizer from a local artifact. - - Args: - artifact: Path to a bundled tokenizer file such as ``tokenizer.json``. - tokenizer_class: Allowlisted Transformers tokenizer class name. - tokenizer_options: Checkpoint-derived tokenizer construction options. - max_token_len: Maximum encoded prompt length. - """ - - def __init__( - self, - artifact: str, - tokenizer_class: str, - tokenizer_options: dict[str, Any] | None = None, - max_token_len: int = 512, - ) -> None: - """Initialize a tokenizer from a bundled local artifact. - - Raises: - FileNotFoundError: If the tokenizer artifact does not exist. - ImportError: If Transformers is not installed. - ValueError: If the requested tokenizer class is not supported. - """ - super().__init__() - artifact_path = Path(artifact) - if not artifact_path.is_file(): - msg = f"Tokenizer artifact does not exist: {artifact_path}" - raise FileNotFoundError(msg) - if tokenizer_class not in _SUPPORTED_TOKENIZER_CLASSES: - msg = f"Unsupported asset tokenizer class: {tokenizer_class!r}" - raise ValueError(msg) - - try: - import transformers # ruff: ignore[PLC0415] - except ImportError as exc: - msg = "Tokenizer requires transformers. Install with: pip install transformers" - raise ImportError(msg) from exc - - tokenizer_type = getattr(transformers, tokenizer_class) - self._tokenizer = tokenizer_type.from_pretrained( - artifact_path.parent, - local_files_only=True, - **(tokenizer_options or {}), - ) - self._max_token_len = max_token_len - - def __call__(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - """Tokenize tasks and replace them with token IDs and masks. - - Returns: - Input values with tasks replaced by token IDs and attention masks. - - Raises: - TypeError: If the task value is not a list. - """ - batch_tasks = inputs[TASK] - if not isinstance(batch_tasks, list): - msg = f"Expected TASK to be a list of strings, got {type(batch_tasks)}" - raise TypeError(msg) - - outputs = dict(inputs) - outputs.pop(TASK) - encoded_tokens = self._tokenizer( - batch_tasks, - max_length=self._max_token_len, - truncation=True, - padding="max_length", - return_tensors="np", - ) - outputs[TOKENIZED_PROMPT] = encoded_tokens["input_ids"] - outputs[TOKENIZED_PROMPT_MASK] = encoded_tokens["attention_mask"].astype(np.bool_) - return outputs - - def __repr__(self) -> str: - """Return string representation of the preprocessor.""" - return f"{self.__class__.__name__}(tokenizer={self._tokenizer.name_or_path!r})" diff --git a/tests/unit/inference/preprocessors/test_asset_tokenizer.py b/tests/unit/inference/preprocessors/test_asset_tokenizer.py deleted file mode 100644 index c5f2353b..00000000 --- a/tests/unit/inference/preprocessors/test_asset_tokenizer.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (C) 2026 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from pathlib import Path -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest - -from physicalai.inference.component_factory import instantiate_component, resolve_artifact -from physicalai.inference.constants import TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK -from physicalai.inference.manifest import ComponentSpec -from physicalai.inference.preprocessors import AssetTokenizer - - -def _mock_tokenizer() -> MagicMock: - tokenizer = MagicMock() - tokenizer.name_or_path = "local-tokenizer" - - def _encode(tasks, **kwargs): - length = kwargs["max_length"] - return { - "input_ids": np.ones((len(tasks), length), dtype=np.int64), - "attention_mask": np.ones((len(tasks), length), dtype=np.int64), - } - - tokenizer.side_effect = _encode - return tokenizer - - -def test_loads_allowlisted_tokenizer_with_dynamic_options(tmp_path: Path) -> None: - tokenizer_file = tmp_path / "tokenizer.json" - tokenizer_file.write_text("{}", encoding="utf-8") - transformers = MagicMock() - tokenizer = _mock_tokenizer() - transformers.Qwen2Tokenizer.from_pretrained.return_value = tokenizer - options = { - "bos_token": "<|im_end|>", - "extra_special_tokens": ["", "<|image|>"], - "model_max_length": 1010000, - } - - with patch.dict("sys.modules", {"transformers": transformers}): - preprocessor = AssetTokenizer( - artifact=str(tokenizer_file), - tokenizer_class="Qwen2Tokenizer", - tokenizer_options=options, - max_token_len=4, - ) - - transformers.Qwen2Tokenizer.from_pretrained.assert_called_once_with( - tmp_path, - local_files_only=True, - **options, - ) - result = preprocessor({TASK: ["pick up the block"]}) - assert result[TOKENIZED_PROMPT].shape == (1, 4) - assert result[TOKENIZED_PROMPT_MASK].dtype == np.bool_ - assert TASK not in result - - -def test_rejects_unsupported_tokenizer_class(tmp_path: Path) -> None: - tokenizer_file = tmp_path / "tokenizer.json" - tokenizer_file.write_text("{}", encoding="utf-8") - - with pytest.raises(ValueError, match="Unsupported asset tokenizer class"): - AssetTokenizer(artifact=str(tokenizer_file), tokenizer_class="ArbitraryTokenizer") - - -def test_manifest_resolves_flat_artifact(tmp_path: Path) -> None: - tokenizer_file = tmp_path / "tokenizer.json" - tokenizer_file.write_text("{}", encoding="utf-8") - transformers = MagicMock() - transformers.Qwen2Tokenizer.from_pretrained.return_value = _mock_tokenizer() - spec = resolve_artifact( - ComponentSpec( - type="asset_tokenizer", - artifact="tokenizer.json", - tokenizer_class="Qwen2Tokenizer", - tokenizer_options={"extra_special_tokens": ["<|image|>"]}, - ), - tmp_path, - ) - - with patch.dict("sys.modules", {"transformers": transformers}): - component = instantiate_component(spec) - - assert isinstance(component, AssetTokenizer) - assert spec.flat_params["artifact"] == str(tokenizer_file) \ No newline at end of file From 80bb41650dad8101fb48045cd31b5bba390aef55 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:20:26 +0100 Subject: [PATCH 09/22] molmo pre / post processors --- .../inference/preprocessors/molmoact2.py | 1 - .../inference/preprocessors/test_molmoact2.py | 169 +++++++----------- 2 files changed, 62 insertions(+), 108 deletions(-) diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 084d4bc4..01d25791 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -442,5 +442,4 @@ def _token_type_ids(self, ids: np.ndarray, mask: np.ndarray) -> np.ndarray | Non return None return (np.isin(ids, self._image_token_ids) & mask.astype(np.bool_)).astype(np.int64) - __all__ = ["MolmoAct2ModelInputs", "MolmoAct2Preprocessor"] diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 0ed33c80..be43aaaa 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -12,7 +12,7 @@ from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor -def _raw_preprocessor(**kwargs) -> MolmoAct2Preprocessor: +def _prepare(**kwargs) -> MolmoAct2Preprocessor: return MolmoAct2Preprocessor( image_keys=["top", "wrist"], image_size=(28, 28), @@ -23,37 +23,41 @@ def _raw_preprocessor(**kwargs) -> MolmoAct2Preprocessor: ) -def _model_inputs(**kwargs) -> MolmoAct2ModelInputs: - return MolmoAct2ModelInputs( - 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_size=(28, 28), - patch_size=14, - pooling_size=(2, 2), - image_token_ids=[10, 11, 12, 13], - **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 = _raw_preprocessor() - inputs = { - STATE: np.array([[-1.0, 1.0]], dtype=np.float32), - TASK: ["Task: Pick up."], - f"{IMAGES}.wrist": np.full((1, 3, 28, 28), 255, dtype=np.uint8), - f"{IMAGES}.top": np.zeros((1, 3, 28, 28), dtype=np.uint8), - } - - result = processor(inputs) + 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 @@ -71,57 +75,19 @@ def test_applies_masked_normalization_and_joint_transform(self) -> None: joint_signs=[1.0, -1.0], joint_offsets=[0.0, 2.0], ) - result = processor( - { - STATE: np.array([[1.0, 1.0]], dtype=np.float32), - TASK: "move", - IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8), - }, - ) + result = processor({ + STATE: np.array([[1.0, 1.0]], dtype=np.float32), + TASK: "move", + IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8), + }) assert "" in result[TASK][0] - def test_rejects_missing_state(self) -> None: - with pytest.raises(ValueError, match="state"): - _raw_preprocessor()({TASK: ["move"], IMAGES: np.zeros((1, 3, 28, 28), dtype=np.uint8)}) - - def test_registry_alias_instantiates(self) -> None: - processor = instantiate_component(ComponentSpec(type="molmoact2", image_keys=[])) - assert isinstance(processor, MolmoAct2Preprocessor) - - -class TestMolmoAct2ModelInputs: - def test_assembles_model_inputs(self) -> None: - processor = _model_inputs() - result = processor( - { - TOKENIZED_PROMPT: np.array([[99, 5, 99, 0]], dtype=np.int64), - TOKENIZED_PROMPT_MASK: np.array([[1, 1, 1, 0]], dtype=np.bool_), - IMAGES: np.zeros((2, 1, 3, 28, 28), dtype=np.float32), - }, - ) - - assert set(result) == { - "input_ids", - "attention_mask", - "token_type_ids", - "images", - "token_pooling", - "action_dim_is_pad", - } - assert int(result["input_ids"][0, 0]) == 1 - assert result["images"].shape == (1, 2, 4, 588) - assert result["token_pooling"].shape == (1, 2, 4) - np.testing.assert_array_equal(result["action_dim_is_pad"], [[False, False, True, True]]) - assert result["token_type_ids"].sum() > 0 - def test_preserves_masked_tokenizer_padding(self) -> None: - result = _model_inputs()( - { - 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), - }, - ) + 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) @@ -129,36 +95,25 @@ def test_preserves_masked_tokenizer_padding(self) -> None: def test_rejects_placeholder_image_mismatch(self) -> None: with pytest.raises(ValueError, match="placeholders"): - _model_inputs()( - { - TOKENIZED_PROMPT: np.array([[99, 5]], dtype=np.int64), - TOKENIZED_PROMPT_MASK: np.ones((1, 2), dtype=np.bool_), - IMAGES: np.zeros((2, 1, 3, 28, 28), dtype=np.float32), - }, - ) - - def test_registry_alias_instantiates(self) -> None: - spec = ComponentSpec( - type="molmoact2_inputs", - 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, - ) - assert isinstance(instantiate_component(spec), MolmoAct2ModelInputs) + _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_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 TestMolmoAct2ManifestPipeline: def test_processes_observation_and_action(self, monkeypatch) -> None: from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer - class StubTokenizer: + class TransformersTokenizer: name_or_path = "allenai/MolmoAct2" config = type("Config", (), {"revision": "1dbc166cf8765166998eff31ade2eb64c8a40076"})() @@ -169,10 +124,7 @@ def __call__(self, tasks, **kwargs): "attention_mask": np.array([[1, 1, 0, 0]], dtype=np.int64), } - monkeypatch.setattr( - "transformers.AutoTokenizer.from_pretrained", - lambda *args, **kwargs: StubTokenizer(), - ) + monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", lambda *args, **kwargs: TransformersTokenizer()) specs = [ ComponentSpec( type="molmoact2", @@ -209,10 +161,13 @@ def __call__(self, tasks, **kwargs): "task": ["pick up the block"], "images.top": np.zeros((1, 3, 28, 28), dtype=np.uint8), } - for spec in specs: - values = instantiate_component(spec)(values) + processors = [instantiate_component(spec) for spec in specs] + for processor in processors: + values = processor(values) - assert isinstance(instantiate_component(specs[1]), HFTokenizer) + assert isinstance(processors[0], MolmoAct2Preprocessor) + assert isinstance(processors[1], HFTokenizer) + assert isinstance(processors[2], MolmoAct2ModelInputs) 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]] From 1f33726188286a2a53975fe8be29fbfe533d6689 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:47:01 +0100 Subject: [PATCH 10/22] fix: prek lines --- src/physicalai/inference/preprocessors/molmoact2.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 01d25791..084d4bc4 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -442,4 +442,5 @@ def _token_type_ids(self, ids: np.ndarray, mask: np.ndarray) -> np.ndarray | Non return None return (np.isin(ids, self._image_token_ids) & mask.astype(np.bool_)).astype(np.int64) + __all__ = ["MolmoAct2ModelInputs", "MolmoAct2Preprocessor"] From 6818fc86a86e138be56261c106346ff9b73924b0 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 12 Aug 2026 10:43:11 +0100 Subject: [PATCH 11/22] prek: fix --- .../inference/postprocessors/molmoact2.py | 11 +++++ .../inference/preprocessors/molmoact2.py | 48 ++++++++++++++++++- .../preprocessors/molmoact2_inputs.py | 24 +++++++++- .../inference/preprocessors/test_molmoact2.py | 19 +++++--- 4 files changed, 93 insertions(+), 9 deletions(-) diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py index 2def501d..29e9eff8 100644 --- a/src/physicalai/inference/postprocessors/molmoact2.py +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -26,6 +26,17 @@ def __init__( joint_signs: list[float] | None = None, joint_offsets: list[float] | None = None, ) -> None: + """Initialize the MolmoAct2 postprocessor. + + Args: + action_stats: Quantile statistics used to denormalize actions. + adapt_to_so101: Whether to transform actions to the SO-101 joint frame. + joint_signs: Per-joint signs used by the SO-101 transform. + joint_offsets: Per-joint offsets used by the SO-101 transform. + + Raises: + ValueError: If ``joint_signs`` and ``joint_offsets`` have different lengths. + """ signs = joint_signs or [] offsets = joint_offsets or [] if len(signs) != len(offsets): diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 084d4bc4..e66a2537 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -40,7 +40,11 @@ def _joint_transform( *, inverse: bool, ) -> np.ndarray: - """Apply the MolmoAct2 joint-frame transform to leading dimensions.""" + """Apply the MolmoAct2 joint-frame transform to leading dimensions. + + Returns: + A copy of ``values`` transformed between model and robot joint frames. + """ count = min(signs.size, values.shape[-1]) output = np.array(values, copy=True) joints = values[..., :count] @@ -120,6 +124,25 @@ def __init__( joint_signs: list[float] | None = None, joint_offsets: list[float] | None = None, ) -> None: + """Initialize prompt and image preprocessing. + + Args: + image_keys: Ordered camera keys to read from each observation. + state_stats: Quantile statistics used to normalize robot state. + image_size: Output image height and width. + num_state_tokens: Number of discrete tokens available per state dimension. + setup_type: Robot setup text included in the prompt. + control_mode: Robot control mode included in the prompt. + add_setup_tokens: Whether to wrap the setup text in special tokens. + add_control_tokens: Whether to wrap the control text in special tokens. + adapt_to_so101: Whether to transform state from the SO-101 joint frame. + joint_signs: Per-joint signs used by the SO-101 transform. + joint_offsets: Per-joint offsets used by the SO-101 transform. + + Raises: + ValueError: If ``num_state_tokens`` is not positive or the joint transform + lists have different lengths. + """ if num_state_tokens <= 0: msg = f"num_state_tokens must be > 0, got {num_state_tokens}" raise ValueError(msg) @@ -272,6 +295,29 @@ def __init__( use_single_crop_start_token: bool = True, image_token_ids: list[int] | None = None, ) -> None: + """Initialize MolmoAct2 model-input assembly. + + Args: + max_action_dim: Width of the padded action dimension mask. + action_dim: Number of action dimensions used by the environment. + bos_token_id: Beginning-of-sequence token identifier. + pad_token_id: Padding token identifier. + image_placeholder_token_id: Prompt token replaced by an image sequence. + image_start_token_id: Image sequence start token identifier. + image_end_token_id: Image sequence end token identifier. + image_patch_id: Image patch token identifier. + image_col_id: Optional image column separator token identifier. + low_res_image_start_token_id: Optional low-resolution image start token. + image_size: Input image height and width. + patch_size: Height and width of each square image patch. + pooling_size: Height and width of each patch-pooling window. + image_mean: Per-channel image normalization means. + image_std: Per-channel image normalization standard deviations. + image_use_col_tokens: Whether high-resolution rows use column tokens. + use_single_crop_col_tokens: Whether single-crop rows use column tokens. + use_single_crop_start_token: Whether single crops use their configured start token. + image_token_ids: Token identifiers marked as image content. + """ self._max_action_dim = max_action_dim self._action_dim = action_dim self._bos_token_id = bos_token_id diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py index c0f6bc5f..67cbcca6 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -114,7 +114,20 @@ def expand_image_placeholders( 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.""" + """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) @@ -169,6 +182,9 @@ def build_batched_images( Returns: ``(images, token_pooling)`` of shapes ``(N, max_crops, n_patches, pixels)`` and ``(N, max_pooled, pool_area)``. + + Raises: + ValueError: If image-end token and image-grid counts differ. """ counts = (input_ids == int(config.image_end_token_id)).sum(1) # images per example num_images = int(image_grids.shape[0]) @@ -231,7 +247,11 @@ def build_batched_images( def default_action_dim_is_pad(config: MolmoAct2InputConfig, *, batch_size: int) -> np.ndarray: - """Mark action dimensions beyond the environment action dim as padding.""" + """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 diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index be43aaaa..8035d83b 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -10,6 +10,7 @@ from physicalai.inference.manifest import ComponentSpec from physicalai.inference.component_factory import instantiate_component from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.postprocessors import MolmoAct2Postprocessor def _prepare(**kwargs) -> MolmoAct2Preprocessor: @@ -161,13 +162,18 @@ def __call__(self, tasks, **kwargs): "task": ["pick up the block"], "images.top": np.zeros((1, 3, 28, 28), dtype=np.uint8), } - processors = [instantiate_component(spec) for spec in specs] - for processor in processors: - values = processor(values) + 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) + values = model_inputs(values) - assert isinstance(processors[0], MolmoAct2Preprocessor) - assert isinstance(processors[1], HFTokenizer) - assert isinstance(processors[2], MolmoAct2ModelInputs) 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]] @@ -178,5 +184,6 @@ def __call__(self, tasks, **kwargs): 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)) From 7eb2ffb0908a2af83968e024e54b491edbae9639 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:29:17 +0100 Subject: [PATCH 12/22] restore: restore changed files --- pyproject.toml | 2 +- src/physicalai/inference/model.py | 23 +---------------------- uv.lock | 2 +- 3 files changed, 3 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 88a8eb38..3719419c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] dependencies = [ - "numpy", + "numpy>=1.24", "loguru>=0.7", # target PyTurboJPEG 1.x for libjpeg-turbo 2.x compatibility "PyTurboJPEG<2; sys_platform == 'linux'", diff --git a/src/physicalai/inference/model.py b/src/physicalai/inference/model.py index 937fd915..55ceed28 100644 --- a/src/physicalai/inference/model.py +++ b/src/physicalai/inference/model.py @@ -301,30 +301,9 @@ def select_action(self, observation: dict[str, np.ndarray]) -> np.ndarray: >>> next_obs, reward, done = env.step(action) """ if not self._action_buffer: - action_chunk = self.predict_action_chunk(observation) - self._action_buffer.extend(action_chunk[: self._effective_chunk_size()]) + self._action_buffer.extend(self.predict_action_chunk(observation)) return self._action_buffer.popleft() - def _effective_chunk_size(self) -> int: - """Return the number of actions to queue per model invocation. - - Preference order: - 1. Runner-declared ``chunk_size`` from the manifest. - 2. ACTION output feature leading dimension (when declared as ``(T, D)``). - 3. Fallback to 1. - """ - runner_chunk = int(self.chunk_size) - if runner_chunk > 1: - return runner_chunk - - for feature in self.output_features: - if feature.name == ACTION and len(feature.shape) >= 2: - action_chunk = int(feature.shape[0]) - if action_chunk > 0: - return action_chunk - - return 1 - def predict_action_chunk(self, observation: dict[str, np.ndarray]) -> np.ndarray: """Predict a chunk of actions for the given observation. diff --git a/uv.lock b/uv.lock index a508be72..8b928314 100644 --- a/uv.lock +++ b/uv.lock @@ -1834,7 +1834,7 @@ requires-dist = [ { name = "msgpack", marker = "extra == 'transport'", specifier = "==1.2.1" }, { name = "mypy", marker = "extra == 'tests'" }, { name = "num2words", specifier = ">=0.5.14,<0.6.0" }, - { name = "numpy" }, + { name = "numpy", specifier = ">=1.24" }, { name = "onnxruntime" }, { name = "opencv-python-headless" }, { name = "opencv-python-headless", marker = "extra == 'basler'" }, From 8b747d536a8d38dc167519fdf95f93fcc949ea30 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:51:57 +0100 Subject: [PATCH 13/22] ruff: refactor for ruff --- .../inference/preprocessors/molmoact2.py | 172 +++++++++--------- .../preprocessors/molmoact2_image.py | 132 ++++++++++---- .../preprocessors/molmoact2_inputs.py | 153 +++++++++++----- 3 files changed, 291 insertions(+), 166 deletions(-) diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index e66a2537..44f20b65 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -6,6 +6,7 @@ from __future__ import annotations import re +from dataclasses import dataclass from typing import Any import cv2 @@ -269,74 +270,73 @@ def _resize_image(self, image: np.ndarray) -> np.ndarray: return np.stack(output, axis=0) +@dataclass(eq=False, repr=False, kw_only=True) class MolmoAct2ModelInputs(Preprocessor): - """Assemble tokenized prompts and packed images into MolmoAct2 model inputs.""" - - def __init__( - self, - *, - 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, - 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_use_col_tokens: bool = True, - use_single_crop_col_tokens: bool = False, - use_single_crop_start_token: bool = True, - image_token_ids: list[int] | None = None, - ) -> None: - """Initialize MolmoAct2 model-input assembly. + """Assemble tokenized prompts and packed images into MolmoAct2 model inputs. + + Args: + max_action_dim: Width of the padded action dimension mask. + action_dim: Number of action dimensions used by the environment. + bos_token_id: Beginning-of-sequence token identifier. + pad_token_id: Padding token identifier. + image_placeholder_token_id: Prompt token replaced by an image sequence. + image_start_token_id: Image sequence start token identifier. + image_end_token_id: Image sequence end token identifier. + image_patch_id: Image patch token identifier. + image_col_id: Optional image column separator token identifier. + low_res_image_start_token_id: Optional low-resolution image start token. + image_size: Input image height and width. + patch_size: Height and width of each square image patch. + pooling_size: Height and width of each patch-pooling window. + image_mean: Per-channel image normalization means. + image_std: Per-channel image normalization standard deviations. + image_use_col_tokens: Whether high-resolution rows use column tokens. + use_single_crop_col_tokens: Whether single-crop rows use column tokens. + use_single_crop_start_token: Whether single crops use their configured start token. + image_token_ids: Token identifiers marked as image content. + """ - Args: - max_action_dim: Width of the padded action dimension mask. - action_dim: Number of action dimensions used by the environment. - bos_token_id: Beginning-of-sequence token identifier. - pad_token_id: Padding token identifier. - image_placeholder_token_id: Prompt token replaced by an image sequence. - image_start_token_id: Image sequence start token identifier. - image_end_token_id: Image sequence end token identifier. - image_patch_id: Image patch token identifier. - image_col_id: Optional image column separator token identifier. - low_res_image_start_token_id: Optional low-resolution image start token. - image_size: Input image height and width. - patch_size: Height and width of each square image patch. - pooling_size: Height and width of each patch-pooling window. - image_mean: Per-channel image normalization means. - image_std: Per-channel image normalization standard deviations. - image_use_col_tokens: Whether high-resolution rows use column tokens. - use_single_crop_col_tokens: Whether single-crop rows use column tokens. - use_single_crop_start_token: Whether single crops use their configured start token. - image_token_ids: Token identifiers marked as image content. - """ - self._max_action_dim = max_action_dim - self._action_dim = action_dim - self._bos_token_id = bos_token_id - self._pad_token_id = pad_token_id - self._placeholder_id = image_placeholder_token_id - self._image_start_id = image_start_token_id - self._image_end_id = image_end_token_id - self._image_patch_id = image_patch_id - self._image_col_id = image_col_id - self._low_res_start_id = low_res_image_start_token_id or image_start_token_id - self._height, self._width = image_size - self._patch_size = patch_size - self._pool_h, self._pool_w = pooling_size - self._mean = np.asarray(image_mean or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) - self._std = np.asarray(image_std or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) - self._image_use_col_tokens = image_use_col_tokens - self._use_single_crop_col_tokens = use_single_crop_col_tokens - self._use_single_crop_start_token = use_single_crop_start_token - self._image_token_ids = np.asarray(image_token_ids or [], dtype=np.int64) + 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 + 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_use_col_tokens: bool = True + use_single_crop_col_tokens: bool = False + use_single_crop_start_token: bool = True + image_token_ids: list[int] | None = None + + def __post_init__(self) -> None: + """Initialize private model-input assembly state.""" + self._max_action_dim = self.max_action_dim + self._action_dim = self.action_dim + self._bos_token_id = self.bos_token_id + self._pad_token_id = self.pad_token_id + self._placeholder_id = self.image_placeholder_token_id + self._image_start_id = self.image_start_token_id + self._image_end_id = self.image_end_token_id + self._image_patch_id = self.image_patch_id + self._image_col_id = self.image_col_id + self._low_res_start_id = self.low_res_image_start_token_id or self.image_start_token_id + self._height, self._width = self.image_size + self._patch_size = self.patch_size + self._pool_h, self._pool_w = self.pooling_size + self._mean = np.asarray(self.image_mean or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) + self._std = np.asarray(self.image_std or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) + self._image_use_col_tokens = self.image_use_col_tokens + self._use_single_crop_col_tokens = self.use_single_crop_col_tokens + self._use_single_crop_start_token = self.use_single_crop_start_token + self._image_token_ids = np.asarray(self.image_token_ids or [], dtype=np.int64) self._pooling, self._pooled_h, self._pooled_w = self._pooling_indices() @override @@ -345,6 +345,26 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: attention_mask = np.asarray(inputs[TOKENIZED_PROMPT_MASK], dtype=np.int64) input_ids, attention_mask = self._insert_bos(input_ids, attention_mask) + pixel_values, grids, num_images, batch_size = self._prepare_images(inputs) + input_ids, attention_mask = self._expand_placeholders(input_ids, attention_mask, grids) + token_type_ids = self._token_type_ids(input_ids, attention_mask) + batched_images = pixel_values.reshape(batch_size, num_images, pixel_values.shape[1], pixel_values.shape[2]) + token_pooling = self._build_token_pooling(pixel_values, num_images=num_images, batch_size=batch_size) + action_dim_is_pad = np.ones((batch_size, self._max_action_dim), dtype=np.bool_) + action_dim_is_pad[:, : self._action_dim] = False + + outputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + IMAGES: batched_images.astype(np.float32), + "token_pooling": token_pooling.astype(np.int64), + "action_dim_is_pad": action_dim_is_pad, + } + if token_type_ids is not None: + outputs["token_type_ids"] = token_type_ids + return outputs + + def _prepare_images(self, inputs: dict[str, Any]) -> tuple[np.ndarray, np.ndarray, int, int]: 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}" @@ -359,29 +379,15 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: grids = np.tile( np.array([[self._pooled_h, self._pooled_w, 0, 0]], dtype=np.int64), (batch_size * num_images, 1) ) - input_ids, attention_mask = self._expand_placeholders(input_ids, attention_mask, grids) - token_type_ids = self._token_type_ids(input_ids, attention_mask) - batched_images = pixel_values.reshape(batch_size, num_images, pixel_values.shape[1], pixel_values.shape[2]) + return pixel_values, grids, num_images, batch_size + def _build_token_pooling(self, pixel_values: np.ndarray, *, num_images: int, batch_size: int) -> np.ndarray: pooling = [] patches_per_image = pixel_values.shape[1] for image_index in range(num_images): block = np.where(self._pooling >= 0, self._pooling + image_index * patches_per_image, self._pooling) pooling.append(block) - token_pooling = np.tile(np.concatenate(pooling, axis=0)[None, ...], (batch_size, 1, 1)) - action_dim_is_pad = np.ones((batch_size, self._max_action_dim), dtype=np.bool_) - action_dim_is_pad[:, : self._action_dim] = False - - outputs = { - "input_ids": input_ids, - "attention_mask": attention_mask, - IMAGES: batched_images.astype(np.float32), - "token_pooling": token_pooling.astype(np.int64), - "action_dim_is_pad": action_dim_is_pad, - } - if token_type_ids is not None: - outputs["token_type_ids"] = token_type_ids - return outputs + return np.tile(np.concatenate(pooling, axis=0)[None, ...], (batch_size, 1, 1)) def _insert_bos(self, ids: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]: rows: list[np.ndarray] = [] diff --git a/src/physicalai/inference/preprocessors/molmoact2_image.py b/src/physicalai/inference/preprocessors/molmoact2_image.py index bf340c9f..ddbec157 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_image.py +++ b/src/physicalai/inference/preprocessors/molmoact2_image.py @@ -5,9 +5,13 @@ from __future__ import annotations +from dataclasses import dataclass + import cv2 import numpy as np +_GRAYSCALE_NDIM = 2 + def _normalize_image(image: np.ndarray, image_mean: list[float], image_std: list[float]) -> np.ndarray: if np.allclose(image_mean, [0.5, 0.5, 0.5]) and np.allclose(image_std, [0.5, 0.5, 0.5]): @@ -21,7 +25,7 @@ def _normalize_image(image: np.ndarray, image_mean: list[float], image_std: list def _resize_image(image: np.ndarray, desired_output_size: list[int]) -> np.ndarray: height, width = int(desired_output_size[0]), int(desired_output_size[1]) resized = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR) - if resized.ndim == 2: + if resized.ndim == _GRAYSCALE_NDIM: resized = resized[:, :, None] if np.issubdtype(image.dtype, np.floating): @@ -67,6 +71,65 @@ def _build_resized_image( return resized, resize_idx +@dataclass(frozen=True) +class _CropGeometry: + left_margin: int + right_margin: int + total_margin_pixels: int + crop_window_size: int + crop_patch_w: int + crop_patch_h: int + crop_size: int + + +def _crop_geometry( + overlap_margins: list[int], base_image_input_size: list[int], image_patch_size: int +) -> _CropGeometry: + left_margin, right_margin = overlap_margins + crop_patches = base_image_input_size[0] // image_patch_size + crop_window_patches = crop_patches - (right_margin + left_margin) + return _CropGeometry( + left_margin=left_margin, + right_margin=right_margin, + total_margin_pixels=image_patch_size * (right_margin + left_margin), + crop_window_size=crop_window_patches * image_patch_size, + crop_patch_w=base_image_input_size[1] // image_patch_size, + crop_patch_h=base_image_input_size[0] // image_patch_size, + crop_size=base_image_input_size[0], + ) + + +def _fill_overlapping_crops( + src: np.ndarray, tiling: np.ndarray, geometry: _CropGeometry +) -> tuple[np.ndarray, np.ndarray]: + n_crops = int(tiling[0] * tiling[1]) + crop_arr = np.zeros([n_crops, geometry.crop_size, geometry.crop_size, 3], dtype=src.dtype) + patch_idx_arr = np.zeros([n_crops, geometry.crop_patch_h, geometry.crop_patch_w], dtype=np.int32) + + on_crop = 0 + for i in range(int(tiling[0])): + y0 = i * geometry.crop_window_size + for j in range(int(tiling[1])): + x0 = j * geometry.crop_window_size + crop_arr[on_crop] = src[y0 : y0 + geometry.crop_size, x0 : x0 + geometry.crop_size] + patch_idx = np.arange(geometry.crop_patch_w * geometry.crop_patch_h).reshape( + geometry.crop_patch_h, geometry.crop_patch_w + ) + patch_idx += on_crop * geometry.crop_patch_h * geometry.crop_patch_w + + if i != 0: + patch_idx[: geometry.left_margin, :] = -1 + if j != 0: + patch_idx[:, : geometry.left_margin] = -1 + if i != int(tiling[0]) - 1: + patch_idx[-geometry.right_margin :, :] = -1 + if j != int(tiling[1]) - 1: + patch_idx[:, -geometry.right_margin :] = -1 + patch_idx_arr[on_crop] = patch_idx + on_crop += 1 + return crop_arr, patch_idx_arr + + def _build_overlapping_crops( image: np.ndarray, max_crops: int, @@ -76,55 +139,27 @@ def _build_overlapping_crops( image_std: list[float], image_patch_size: int, ) -> tuple[np.ndarray, np.ndarray]: - left_margin, right_margin = overlap_margins - total_margin_pixels = image_patch_size * (right_margin + left_margin) - crop_patches = base_image_input_size[0] // image_patch_size - crop_window_patches = crop_patches - (right_margin + left_margin) - crop_window_size = crop_window_patches * image_patch_size - crop_patch_w = base_image_input_size[1] // image_patch_size - crop_patch_h = base_image_input_size[0] // image_patch_size - + geometry = _crop_geometry(overlap_margins, base_image_input_size, image_patch_size) original_image_h, original_image_w = image.shape[:2] - crop_size = base_image_input_size[0] tiling = _select_tiling( - original_image_h - total_margin_pixels, - original_image_w - total_margin_pixels, - crop_window_size, + original_image_h - geometry.total_margin_pixels, + original_image_w - geometry.total_margin_pixels, + geometry.crop_window_size, max_crops, ) src = _resize_image( image, - [tiling[0] * crop_window_size + total_margin_pixels, tiling[1] * crop_window_size + total_margin_pixels], + [ + tiling[0] * geometry.crop_window_size + geometry.total_margin_pixels, + tiling[1] * geometry.crop_window_size + geometry.total_margin_pixels, + ], ) src = _normalize_image(src, image_mean, image_std) - n_crops = int(tiling[0] * tiling[1]) - crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype) - patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32) - - on_crop = 0 - for i in range(int(tiling[0])): - y0 = i * crop_window_size - for j in range(int(tiling[1])): - x0 = j * crop_window_size - crop_arr[on_crop] = src[y0 : y0 + crop_size, x0 : x0 + crop_size] - patch_idx = np.arange(crop_patch_w * crop_patch_h).reshape(crop_patch_h, crop_patch_w) - patch_idx += on_crop * crop_patch_h * crop_patch_w - - if i != 0: - patch_idx[:left_margin, :] = -1 - if j != 0: - patch_idx[:, :left_margin] = -1 - if i != int(tiling[0]) - 1: - patch_idx[-right_margin:, :] = -1 - if j != int(tiling[1]) - 1: - patch_idx[:, -right_margin:] = -1 - patch_idx_arr[on_crop] = patch_idx - on_crop += 1 - - patch_idx_arr = patch_idx_arr.reshape(int(tiling[0]), int(tiling[1]), crop_patch_h, crop_patch_w) + crop_arr, patch_idx_arr = _fill_overlapping_crops(src, tiling, geometry) + patch_idx_arr = patch_idx_arr.reshape(int(tiling[0]), int(tiling[1]), geometry.crop_patch_h, geometry.crop_patch_w) patch_idx_arr = patch_idx_arr.transpose(0, 2, 1, 3).reshape(-1) patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape( src.shape[0] // image_patch_size, @@ -242,6 +277,7 @@ class MolmoAct2ImageProcessor: def __init__( self, + *, size: dict[str, int] | None = None, image_mean: list[float] | None = None, image_std: list[float] | None = None, @@ -252,6 +288,19 @@ def __init__( patch_size: int = 14, pooling_size: list[int] | None = None, ) -> None: + """Initialize image processing options. + + Args: + size: Output image height and width. + image_mean: Per-channel image normalization means. + image_std: Per-channel image normalization standard deviations. + do_convert_rgb: Whether to convert images to RGB. + max_crops: Maximum number of overlapping crops. + overlap_margins: Left and right overlap margins in patches. + crop_mode: Image crop strategy. + patch_size: Height and width of each square image patch. + pooling_size: Width and height of each patch-pooling window. + """ self.size = size if size is not None else {"height": 378, "width": 378} self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5] self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5] @@ -263,6 +312,11 @@ def __init__( self.pooling_size = pooling_size if pooling_size is not None else [2, 2] def __call__(self, images_bchw: np.ndarray) -> dict[str, np.ndarray]: + """Convert BCHW images into patches and pooling metadata. + + Returns: + Model-ready pixel values, pooling indices, image grids, and crop counts. + """ image_list = _to_hwc_uint8(images_bchw) patch_batches: list[np.ndarray] = [] pooling_batches: list[np.ndarray] = [] diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py index 67cbcca6..258c3348 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -38,6 +38,7 @@ class MolmoAct2InputConfig: _image_token_ids: list[int] = field(default_factory=list, init=False, repr=False) def __post_init__(self) -> None: + """Collect the configured image token identifiers.""" ids = [ self.image_patch_id, self.image_col_id, @@ -52,12 +53,20 @@ def __post_init__(self) -> None: @property def image_token_ids(self) -> list[int]: - """Token ids that mark image content (for token type ids).""" + """Token ids that mark image content (for token type ids). + + Returns: + Configured image token identifiers. + """ return self._image_token_ids 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.""" + """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) @@ -98,7 +107,11 @@ def make_rows(num_rows: int, num_cols: int, *, use_col: bool) -> list[int]: 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.""" + """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 @@ -165,6 +178,85 @@ def expand_image_placeholders( 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), counts) + 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, @@ -186,61 +278,34 @@ def build_batched_images( Raises: ValueError: If image-end token and image-grid counts differ. """ - counts = (input_ids == int(config.image_end_token_id)).sum(1) # images per example + counts = (input_ids == int(config.image_end_token_id)).sum(1) num_images = int(image_grids.shape[0]) if int(counts.sum()) != num_images: msg = f"image_end tokens ({int(counts.sum())}) do not match image grids ({num_images})." raise ValueError(msg) - num_examples = counts.shape[0] - n_crops, n_patches, pixels_per_patch = pixel_values.shape - del n_crops - - 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), counts) - 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) - patches_per_image = image_num_crops.astype(np.int64) * n_patches - - max_crops = int(crops_per_example.max()) if num_examples > 0 else 0 - images = np.full( - (num_examples, max_crops, n_patches, pixels_per_patch), - -1.0, - dtype=pixel_values.dtype, - ) - max_pooled = int(pooled_per_example.max()) if num_examples > 0 else 0 - token_pooling = np.full( - (num_examples, max_pooled, image_token_pooling.shape[-1]), - -1, - dtype=image_token_pooling.dtype, - ) + 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(num_examples): - num_example_images = int(counts[example_idx]) - num_example_crops = int(crops_per_example[example_idx]) + 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 = image_token_pooling[ - pooled_offset : pooled_offset + int(pooled_per_example[example_idx]) - ].copy() - patch_offset = 0 - row = 0 - for local_image in range(num_example_images): - num_pooled = int(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(patches_per_image[image_offset + local_image]) - row += num_pooled + 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(pooled_per_example[example_idx]) + pooled_offset += int(layout.pooled_per_example[example_idx]) image_offset += num_example_images return images, token_pooling From 6bd331d11a6b25efb4d5b7bbf8ccda50712158d0 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:56:12 +0100 Subject: [PATCH 14/22] fix: accept B, H, W, C input --- .../inference/preprocessors/molmoact2.py | 15 ++++++++++++--- .../inference/preprocessors/test_molmoact2.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 44f20b65..785e4c95 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -224,14 +224,23 @@ def _extract_images(self, inputs: dict[str, Any], *, batch_size: int) -> list[np if not images: msg = "MolmoAct2 requires at least one image input" raise ValueError(msg) + normalized_images: list[np.ndarray] = [] for image in images: - if image.ndim != _IMAGE_NDIM or image.shape[1] != _NUM_CHANNELS: - msg = f"Expected BCHW image with 3 channels, got {image.shape}" + 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] == _NUM_CHANNELS: + normalized_image = image + elif image.shape[-1] == _NUM_CHANNELS: + normalized_image = np.transpose(image, (0, 3, 1, 2)) + else: + msg = f"Expected BCHW or BHWC image with 3 channels, got {image.shape}" raise ValueError(msg) if image.shape[0] != batch_size: msg = f"Image batch size mismatch: expected {batch_size}, got {image.shape[0]}" raise ValueError(msg) - return images + normalized_images.append(normalized_image) + return normalized_images @staticmethod def _extract_tasks(inputs: dict[str, Any], *, batch_size: int) -> list[str]: diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 8035d83b..79df55ee 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -67,6 +67,17 @@ def test_builds_prompt_and_packs_ordered_cameras(self) -> None: assert "The task is to pick up." in result[TASK][0] assert "" in result[TASK][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_and_joint_transform(self) -> None: processor = MolmoAct2Preprocessor( image_keys=[], From 3a59746e4ff5d2f65a85e57e8b8ad3c9178b914a Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:48:19 +0200 Subject: [PATCH 15/22] refactor: simplify and slim, seperate into individual fiels --- .../inference/postprocessors/molmoact2.py | 54 +- .../inference/preprocessors/__init__.py | 3 +- .../inference/preprocessors/molmoact2.py | 570 ++++++------------ .../preprocessors/molmoact2_image.py | 406 +++---------- .../preprocessors/molmoact2_inputs.py | 154 ++++- .../postprocessors/test_molmoact2.py | 11 + .../inference/preprocessors/test_molmoact2.py | 42 ++ 7 files changed, 473 insertions(+), 767 deletions(-) diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py index 29e9eff8..c5bfa018 100644 --- a/src/physicalai/inference/postprocessors/molmoact2.py +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -1,7 +1,7 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -"""NumPy postprocessor for MolmoAct2 exported models.""" +"""NumPy postprocessing for MolmoAct2 inference.""" from __future__ import annotations @@ -13,59 +13,53 @@ 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 JointFrameTransform, normalization_stats class MolmoAct2Postprocessor(Postprocessor): - """Clamp, denormalize, and optionally transform MolmoAct2 actions.""" + """Clamp, denormalize, and optionally convert actions to robot frame.""" def __init__( self, *, action_stats: dict[str, Any] | None = None, + normalization_mode: str = "QUANTILES", adapt_to_so101: bool = False, joint_signs: list[float] | None = None, joint_offsets: list[float] | None = None, ) -> None: - """Initialize the MolmoAct2 postprocessor. - - Args: - action_stats: Quantile statistics used to denormalize actions. - adapt_to_so101: Whether to transform actions to the SO-101 joint frame. - joint_signs: Per-joint signs used by the SO-101 transform. - joint_offsets: Per-joint offsets used by the SO-101 transform. - - Raises: - ValueError: If ``joint_signs`` and ``joint_offsets`` have different lengths. - """ - signs = joint_signs or [] - offsets = joint_offsets or [] - if len(signs) != len(offsets): - msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" - raise ValueError(msg) - self._adapt_to_so101 = adapt_to_so101 - self._joint_signs = np.asarray(signs, dtype=np.float32) - self._joint_offsets = np.asarray(offsets, dtype=np.float32) - self._denormalizer = ( - StatsDenormalizer(stats={ACTION: action_stats}, mode="quantiles", features=[ACTION]) + """Store action postprocessing settings.""" + self.denormalizer = ( + StatsDenormalizer( + stats={ACTION: normalization_stats(action_stats)}, + mode=normalization_mode.lower(), + features=[ACTION], + ) if action_stats else None ) + self.joint_transform = JointFrameTransform(joint_signs, joint_offsets) if adapt_to_so101 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 no action output is present. + """ result = dict(outputs) action = result.get(ACTION, result.get("actions")) if action is None: msg = "MolmoAct2 postprocessor expected an action tensor" raise ValueError(msg) action = np.clip(np.asarray(action), -1.0, 1.0) - if self._denormalizer is not None: - action = self._denormalizer({ACTION: action})[ACTION] - if self._adapt_to_so101: - count = min(self._joint_signs.size, action.shape[-1]) - transformed = np.array(action, copy=True) - transformed[..., :count] = self._joint_signs[:count] * (action[..., :count] - self._joint_offsets[:count]) - action = transformed + if self.denormalizer is not None: + action = self.denormalizer({ACTION: action})[ACTION] + if self.joint_transform is not None: + action = self.joint_transform.apply(action, inverse=True) result.pop("actions", None) result[ACTION] = action return result diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index 2347439f..98134715 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -10,7 +10,8 @@ from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.hf_tokenizer import HFTokenizer from physicalai.inference.preprocessors.lambda_processor import LambdaPreprocessor -from physicalai.inference.preprocessors.molmoact2 import MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.preprocessors.molmoact2 import MolmoAct2Preprocessor +from physicalai.inference.preprocessors.molmoact2_inputs import MolmoAct2ModelInputs from physicalai.inference.preprocessors.new_line import NewLinePreprocessor from physicalai.inference.preprocessors.ov_tokenizer import OVTokenizer from physicalai.inference.preprocessors.pi05 import Pi05Preprocessor diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 785e4c95..75240702 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -1,27 +1,27 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -"""NumPy preprocessors for MolmoAct2 exported models.""" +"""NumPy observation preprocessing for MolmoAct2 inference.""" from __future__ import annotations import re -from dataclasses import dataclass from typing import Any import cv2 import numpy as np from typing_extensions import override -from physicalai.inference.constants import IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK +from physicalai.inference.constants import IMAGES, STATE, TASK from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.stats_normalizer import StatsNormalizer -_STATE_START_TOKEN = "" # noqa: S105 -_STATE_END_TOKEN = "" # noqa: S105 -_STATE_TOKEN_PREFIX = " np.ndarray: - """Apply the MolmoAct2 joint-frame transform to leading dimensions. - - Returns: - A copy of ``values`` transformed between model and robot joint frames. - """ - count = min(signs.size, values.shape[-1]) - output = np.array(values, copy=True) - joints = values[..., :count] - output[..., :count] = ( - signs[:count] * (joints - offsets[:count]) if inverse else signs[:count] * joints + offsets[:count] - ) - return output def _normalize_text(text: str) -> str: - normalized = re.sub(r"\s+", " ", str(text or "")).strip() - if not normalized: - return "" + text = re.sub(r"\s+", " ", str(text or "")).strip() for pattern in _PREFIX_PATTERNS: - normalized = pattern.sub("", normalized, count=1).strip() - return normalized.rstrip(_TRAILING_PUNCTUATION).strip().lower() + text = pattern.sub("", text, count=1).strip() + return text.rstrip(_TRAILING_PUNCTUATION).strip().lower() -def _discrete_state_string(state: np.ndarray, num_state_tokens: int) -> str: - values = np.nan_to_num(np.asarray(state, dtype=np.float32), nan=0.0, posinf=1.0, neginf=-1.0) - values = np.clip(values, -1.0, 1.0) - token_ids = np.rint((values + 1.0) / 2.0 * (num_state_tokens - 1)).astype(np.int64) - payload = "".join(f"{_STATE_TOKEN_PREFIX}{int(token_id)}>" for token_id in token_ids.reshape(-1)) - return f"{_STATE_START_TOKEN}{payload}{_STATE_END_TOKEN}" +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 _wrapped_text(value: str, start: str, end: str, *, enabled: bool) -> str: +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 _robot_prompt( - *, +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, @@ -89,32 +78,68 @@ def _robot_prompt( add_control_tokens: bool, num_images: int, ) -> str: - setup = _wrapped_text(setup_type, "", "", enabled=add_setup_tokens) - control = _wrapped_text(control_mode, "", "", enabled=add_control_tokens) - discrete_state = _discrete_state_string(state, num_state_tokens) + 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}. " + 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|>" - elif num_images > 1: - image_prefix = "".join(f"Image {index + 1}<|image|>" for index in range(num_images)) else: - image_prefix = "" - return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{_ACTION_OUTPUT_TOKEN}" + 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 JointFrameTransform: + """Map leading joints between robot and checkpoint frames.""" + + def __init__( + self, + signs: list[float] | None = None, + offsets: list[float] | None = None, + ) -> None: + """Use Studio SO101 defaults unless compatible overrides are supplied. + + Raises: + ValueError: If signs and offsets have different lengths. + """ + signs = list(SO101_JOINT_SIGNS) if signs is None else signs + offsets = list(SO101_JOINT_OFFSETS) if offsets is None else offsets + if len(signs) != len(offsets): + msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" + raise ValueError(msg) + self.signs = np.asarray(signs, dtype=np.float32) + self.offsets = np.asarray(offsets, dtype=np.float32) + + def apply(self, values: np.ndarray, *, inverse: bool) -> np.ndarray: + """Apply the forward or inverse affine transform. + + Returns: + A transformed copy of the input values. + """ + 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 MolmoAct2Preprocessor(Preprocessor): - """Prepare MolmoAct2 prompts and images before tokenization.""" + """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 = "", @@ -125,50 +150,42 @@ def __init__( joint_signs: list[float] | None = None, joint_offsets: list[float] | None = None, ) -> None: - """Initialize prompt and image preprocessing. - - Args: - image_keys: Ordered camera keys to read from each observation. - state_stats: Quantile statistics used to normalize robot state. - image_size: Output image height and width. - num_state_tokens: Number of discrete tokens available per state dimension. - setup_type: Robot setup text included in the prompt. - control_mode: Robot control mode included in the prompt. - add_setup_tokens: Whether to wrap the setup text in special tokens. - add_control_tokens: Whether to wrap the control text in special tokens. - adapt_to_so101: Whether to transform state from the SO-101 joint frame. - joint_signs: Per-joint signs used by the SO-101 transform. - joint_offsets: Per-joint offsets used by the SO-101 transform. + """Store observation preprocessing settings. Raises: - ValueError: If ``num_state_tokens`` is not positive or the joint transform - lists have different lengths. + ValueError: If no state tokens are available. """ if num_state_tokens <= 0: - msg = f"num_state_tokens must be > 0, got {num_state_tokens}" + msg = f"num_state_tokens must be > 0, got {num_state_tokens}." raise ValueError(msg) - signs = joint_signs or [] - offsets = joint_offsets or [] - if len(signs) != len(offsets): - msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" - 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._adapt_to_so101 = adapt_to_so101 - self._joint_signs = np.asarray(signs, dtype=np.float32) - self._joint_offsets = np.asarray(offsets, dtype=np.float32) - self._normalizer = ( - StatsNormalizer(stats={STATE: state_stats}, mode="quantiles", features=[STATE]) if state_stats else None + 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.joint_transform = JointFrameTransform(joint_signs, joint_offsets) if adapt_to_so101 else None + 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: @@ -177,24 +194,24 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: state = np.asarray(state, dtype=np.float32) if state.ndim == 1: state = state[None, :] - if self._adapt_to_so101: - state = _joint_transform(state, self._joint_signs, self._joint_offsets, inverse=False) - if self._normalizer is not None: - state = self._normalizer({STATE: state})[STATE] + if self.joint_transform is not None: + state = self.joint_transform.apply(state, inverse=False) + if self.normalizer is not None: + state = self.normalizer({STATE: state})[STATE] state = np.clip(state, -1.0, 1.0) - images = self._extract_images(outputs, batch_size=state.shape[0]) - tasks = self._extract_tasks(outputs, batch_size=state.shape[0]) - outputs[IMAGES] = np.stack([self._resize_image(image) for image in images], axis=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] = [ - _robot_prompt( - task=tasks[index], - state=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, + _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]) @@ -203,305 +220,84 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: outputs.pop(f"observation.{STATE}", None) return outputs - def _extract_images(self, inputs: dict[str, Any], *, batch_size: int) -> list[np.ndarray]: - images_value = inputs.get(IMAGES) - images: list[np.ndarray] = [] - if self._image_keys: - for name in self._image_keys: - flat_key = name if name.startswith(f"{IMAGES}.") else f"{IMAGES}.{name}" - if flat_key in inputs: - images.append(np.asarray(inputs[flat_key])) - elif isinstance(images_value, dict) and name.removeprefix(f"{IMAGES}.") in images_value: - images.append(np.asarray(images_value[name.removeprefix(f"{IMAGES}.")])) - elif isinstance(images_value, np.ndarray): - images = [images_value] - elif isinstance(images_value, dict): - images = [np.asarray(value) for key, value in images_value.items() if "is_pad" not in str(key)] - else: - 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) - normalized_images: list[np.ndarray] = [] + 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] == _NUM_CHANNELS: - normalized_image = image - elif image.shape[-1] == _NUM_CHANNELS: - normalized_image = np.transpose(image, (0, 3, 1, 2)) + 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 image.shape[0] != batch_size: - msg = f"Image batch size mismatch: expected {batch_size}, got {image.shape[0]}" + if canonical.shape[0] != batch_size: + msg = f"Image batch size mismatch: expected {batch_size}, got {canonical.shape[0]}" raise ValueError(msg) - normalized_images.append(normalized_image) - return normalized_images + 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): + images = [np.asarray(value) for key, value in container.items() if "is_pad" not in str(key)] + 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 _extract_tasks(inputs: dict[str, Any], *, batch_size: int) -> list[str]: + 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) - if isinstance(source, str): - tasks = [source] * batch_size - else: - tasks = [str(value) for value in np.asarray(source).reshape(-1).tolist()] - if len(tasks) == 1 and batch_size > 1: - tasks *= batch_size + 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(task) for task in tasks] - - def _resize_image(self, image: np.ndarray) -> np.ndarray: - height, width = self._image_size - output: list[np.ndarray] = [] - for sample in image: - if sample.dtype == np.uint8: - pixels = sample - elif np.issubdtype(sample.dtype, np.floating): - float_pixels = sample.astype(np.float32) - if float(np.max(float_pixels)) <= 1.0: - float_pixels *= 255.0 - pixels = np.clip(float_pixels, 0.0, 255.0).astype(np.uint8) + 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: {sample.dtype}" + msg = f"Unsupported image dtype: {image.dtype}" raise ValueError(msg) - hwc = np.transpose(pixels, (1, 2, 0)) - resized = cv2.resize(hwc, (width, height), interpolation=cv2.INTER_LINEAR_EXACT) - output.append(np.transpose(resized, (2, 0, 1)).astype(np.float32) / 255.0) - return np.stack(output, axis=0) - - -@dataclass(eq=False, repr=False, kw_only=True) -class MolmoAct2ModelInputs(Preprocessor): - """Assemble tokenized prompts and packed images into MolmoAct2 model inputs. - - Args: - max_action_dim: Width of the padded action dimension mask. - action_dim: Number of action dimensions used by the environment. - bos_token_id: Beginning-of-sequence token identifier. - pad_token_id: Padding token identifier. - image_placeholder_token_id: Prompt token replaced by an image sequence. - image_start_token_id: Image sequence start token identifier. - image_end_token_id: Image sequence end token identifier. - image_patch_id: Image patch token identifier. - image_col_id: Optional image column separator token identifier. - low_res_image_start_token_id: Optional low-resolution image start token. - image_size: Input image height and width. - patch_size: Height and width of each square image patch. - pooling_size: Height and width of each patch-pooling window. - image_mean: Per-channel image normalization means. - image_std: Per-channel image normalization standard deviations. - image_use_col_tokens: Whether high-resolution rows use column tokens. - use_single_crop_col_tokens: Whether single-crop rows use column tokens. - use_single_crop_start_token: Whether single crops use their configured start token. - image_token_ids: Token identifiers marked as image content. - """ - - 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 - 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_use_col_tokens: bool = True - use_single_crop_col_tokens: bool = False - use_single_crop_start_token: bool = True - image_token_ids: list[int] | None = None - - def __post_init__(self) -> None: - """Initialize private model-input assembly state.""" - self._max_action_dim = self.max_action_dim - self._action_dim = self.action_dim - self._bos_token_id = self.bos_token_id - self._pad_token_id = self.pad_token_id - self._placeholder_id = self.image_placeholder_token_id - self._image_start_id = self.image_start_token_id - self._image_end_id = self.image_end_token_id - self._image_patch_id = self.image_patch_id - self._image_col_id = self.image_col_id - self._low_res_start_id = self.low_res_image_start_token_id or self.image_start_token_id - self._height, self._width = self.image_size - self._patch_size = self.patch_size - self._pool_h, self._pool_w = self.pooling_size - self._mean = np.asarray(self.image_mean or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) - self._std = np.asarray(self.image_std or [0.5, 0.5, 0.5], dtype=np.float32).reshape(1, 3, 1, 1) - self._image_use_col_tokens = self.image_use_col_tokens - self._use_single_crop_col_tokens = self.use_single_crop_col_tokens - self._use_single_crop_start_token = self.use_single_crop_start_token - self._image_token_ids = np.asarray(self.image_token_ids or [], dtype=np.int64) - self._pooling, self._pooled_h, self._pooled_w = self._pooling_indices() - - @override - def __call__(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: - 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) - - pixel_values, grids, num_images, batch_size = self._prepare_images(inputs) - input_ids, attention_mask = self._expand_placeholders(input_ids, attention_mask, grids) - token_type_ids = self._token_type_ids(input_ids, attention_mask) - batched_images = pixel_values.reshape(batch_size, num_images, pixel_values.shape[1], pixel_values.shape[2]) - token_pooling = self._build_token_pooling(pixel_values, num_images=num_images, batch_size=batch_size) - action_dim_is_pad = np.ones((batch_size, self._max_action_dim), dtype=np.bool_) - action_dim_is_pad[:, : self._action_dim] = False - - outputs = { - "input_ids": input_ids, - "attention_mask": attention_mask, - IMAGES: batched_images.astype(np.float32), - "token_pooling": token_pooling.astype(np.int64), - "action_dim_is_pad": action_dim_is_pad, - } - if token_type_ids is not None: - outputs["token_type_ids"] = token_type_ids - return outputs - - def _prepare_images(self, inputs: dict[str, Any]) -> tuple[np.ndarray, np.ndarray, int, int]: - 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 - if channels != _NUM_CHANNELS or (height, width) != (self._height, self._width): - msg = f"Unexpected packed image shape {images.shape}" - raise ValueError(msg) - - flat_images = images.transpose(1, 0, 2, 3, 4).reshape(batch_size * num_images, channels, height, width) - pixel_values = self._patchify((flat_images - self._mean) / self._std) - grids = np.tile( - np.array([[self._pooled_h, self._pooled_w, 0, 0]], dtype=np.int64), (batch_size * num_images, 1) - ) - return pixel_values, grids, num_images, batch_size - - def _build_token_pooling(self, pixel_values: np.ndarray, *, num_images: int, batch_size: int) -> np.ndarray: - pooling = [] - patches_per_image = pixel_values.shape[1] - for image_index in range(num_images): - block = np.where(self._pooling >= 0, self._pooling + image_index * patches_per_image, self._pooling) - pooling.append(block) - return np.tile(np.concatenate(pooling, axis=0)[None, ...], (batch_size, 1, 1)) - - def _insert_bos(self, ids: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - rows: list[np.ndarray] = [] - for row_ids, row_mask in zip(ids, mask, strict=True): - valid_ids = row_ids[row_mask.astype(np.bool_)] - if valid_ids.size == 0 or valid_ids[0] != self._bos_token_id: - valid_ids = np.concatenate((np.array([self._bos_token_id], dtype=ids.dtype), valid_ids)) - rows.append(valid_ids) - 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 - - 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, - ) - pooling = indices.reshape(pooled_h, self._pool_h, pooled_w, self._pool_w) - return pooling.transpose(0, 2, 1, 3).reshape(-1, self._pool_h * self._pool_w), pooled_h, pooled_w - - def _image_sequence(self, grid: np.ndarray) -> list[int]: - resized_h, resized_w, height, width = (int(value) for value in grid) - - def rows(row_count: int, col_count: int, *, use_col: bool) -> list[int]: - row = [self._image_patch_id] * col_count - if use_col and self._image_col_id is not None: - row.append(self._image_col_id) - return row * row_count - - if height == 0 or width == 0: - return [ - self._image_start_id, - *rows(resized_h, resized_w, use_col=self._use_single_crop_col_tokens), - self._image_end_id, - ] - low_start = self._low_res_start_id if self._use_single_crop_start_token else self._image_start_id - return [ - low_start, - *rows(resized_h, resized_w, use_col=self._use_single_crop_col_tokens), - self._image_end_id, - self._image_start_id, - *rows(height, width, use_col=self._image_use_col_tokens), - self._image_end_id, - ] - - def _expand_placeholders( - self, - ids: np.ndarray, - mask: np.ndarray, - grids: np.ndarray, - ) -> tuple[np.ndarray, np.ndarray]: - rows: list[np.ndarray] = [] - expanded_widths: list[int] = [] - grid_index = 0 - for row_ids, row_mask in zip(ids, mask, strict=True): - valid = row_mask.astype(np.bool_) - expanded: list[int] = [] - for token in row_ids[valid]: - if int(token) == self._placeholder_id: - if grid_index >= grids.shape[0]: - msg = "Not enough image grids to expand all <|image|> placeholders" - raise ValueError(msg) - expanded.extend(self._image_sequence(grids[grid_index])) - grid_index += 1 - else: - expanded.append(int(token)) - rows.append(np.asarray(expanded, dtype=ids.dtype)) - expanded_widths.append(len(expanded) + int((~valid).sum())) - if grid_index != grids.shape[0]: - msg = f"Image placeholders ({grid_index}) do not match images ({grids.shape[0]})" - raise ValueError(msg) - width = max(expanded_widths, default=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 - - def _token_type_ids(self, ids: np.ndarray, mask: np.ndarray) -> np.ndarray | None: - if self._image_token_ids.size == 0: - return None - return (np.isin(ids, self._image_token_ids) & mask.astype(np.bool_)).astype(np.int64) - - -__all__ = ["MolmoAct2ModelInputs", "MolmoAct2Preprocessor"] + 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__ = [ + "SO101_JOINT_OFFSETS", + "SO101_JOINT_SIGNS", + "JointFrameTransform", + "MolmoAct2Preprocessor", + "normalization_stats", +] diff --git a/src/physicalai/inference/preprocessors/molmoact2_image.py b/src/physicalai/inference/preprocessors/molmoact2_image.py index ddbec157..920bb67c 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_image.py +++ b/src/physicalai/inference/preprocessors/molmoact2_image.py @@ -1,364 +1,102 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -"""NumPy image preprocessing helpers for MolmoAct2 inference.""" +"""NumPy image patchification for MolmoAct2 inference.""" from __future__ import annotations -from dataclasses import dataclass - -import cv2 import numpy as np -_GRAYSCALE_NDIM = 2 - - -def _normalize_image(image: np.ndarray, image_mean: list[float], image_std: list[float]) -> np.ndarray: - if np.allclose(image_mean, [0.5, 0.5, 0.5]) and np.allclose(image_std, [0.5, 0.5, 0.5]): - return image * np.asarray(2.0, dtype=np.float32) - np.asarray(1.0, dtype=np.float32) - image = image.astype(np.float32) - image -= np.asarray(image_mean, dtype=np.float32)[None, None, :] - image /= np.asarray(image_std, dtype=np.float32)[None, None, :] - return image - - -def _resize_image(image: np.ndarray, desired_output_size: list[int]) -> np.ndarray: - height, width = int(desired_output_size[0]), int(desired_output_size[1]) - resized = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR) - if resized.ndim == _GRAYSCALE_NDIM: - resized = resized[:, :, None] - - if np.issubdtype(image.dtype, np.floating): - resized = np.clip(resized, 0.0, 1.0).astype(np.float32) - else: - resized = resized.astype(np.float32) / 255.0 - return resized - - -def _select_tiling(h: int, w: int, patch_size: int, max_num_crops: int) -> np.ndarray: - tilings: list[tuple[int, int]] = [] - for i in range(1, max_num_crops + 1): - tilings.extend((i, j) for j in range(1, max_num_crops + 1) if i * j <= max_num_crops) - tilings.sort(key=lambda x: (x[0] * x[1], x[0])) - candidate_tilings = np.asarray(tilings, dtype=np.int32) - candidate_resolutions = candidate_tilings * patch_size - - original_size = np.asarray([h, w], dtype=np.float32) - with np.errstate(divide="ignore"): - required_scale = candidate_resolutions.astype(np.float32) / original_size[None, :] - required_scale = np.min(required_scale, axis=-1, keepdims=True) - if np.all(required_scale < 1): - ix = int(np.argmax(required_scale)) - else: - required_scale = np.where(required_scale < 1.0, 1e10, required_scale) - ix = int(np.argmin(required_scale)) - return candidate_tilings[ix] - - -def _build_resized_image( - image: np.ndarray, - base_image_input_size: list[int], - image_mean: list[float], - image_std: list[float], - image_patch_size: int, -) -> tuple[np.ndarray, np.ndarray]: - resized = _resize_image(image, base_image_input_size) - resized = _normalize_image(resized, image_mean, image_std) - resized = resized[None, ...] - crop_patch_w = base_image_input_size[1] // image_patch_size - crop_patch_h = base_image_input_size[0] // image_patch_size - resize_idx = np.arange(crop_patch_w * crop_patch_h).reshape([crop_patch_h, crop_patch_w]) - return resized, resize_idx - - -@dataclass(frozen=True) -class _CropGeometry: - left_margin: int - right_margin: int - total_margin_pixels: int - crop_window_size: int - crop_patch_w: int - crop_patch_h: int - crop_size: int - - -def _crop_geometry( - overlap_margins: list[int], base_image_input_size: list[int], image_patch_size: int -) -> _CropGeometry: - left_margin, right_margin = overlap_margins - crop_patches = base_image_input_size[0] // image_patch_size - crop_window_patches = crop_patches - (right_margin + left_margin) - return _CropGeometry( - left_margin=left_margin, - right_margin=right_margin, - total_margin_pixels=image_patch_size * (right_margin + left_margin), - crop_window_size=crop_window_patches * image_patch_size, - crop_patch_w=base_image_input_size[1] // image_patch_size, - crop_patch_h=base_image_input_size[0] // image_patch_size, - crop_size=base_image_input_size[0], - ) - - -def _fill_overlapping_crops( - src: np.ndarray, tiling: np.ndarray, geometry: _CropGeometry -) -> tuple[np.ndarray, np.ndarray]: - n_crops = int(tiling[0] * tiling[1]) - crop_arr = np.zeros([n_crops, geometry.crop_size, geometry.crop_size, 3], dtype=src.dtype) - patch_idx_arr = np.zeros([n_crops, geometry.crop_patch_h, geometry.crop_patch_w], dtype=np.int32) - - on_crop = 0 - for i in range(int(tiling[0])): - y0 = i * geometry.crop_window_size - for j in range(int(tiling[1])): - x0 = j * geometry.crop_window_size - crop_arr[on_crop] = src[y0 : y0 + geometry.crop_size, x0 : x0 + geometry.crop_size] - patch_idx = np.arange(geometry.crop_patch_w * geometry.crop_patch_h).reshape( - geometry.crop_patch_h, geometry.crop_patch_w - ) - patch_idx += on_crop * geometry.crop_patch_h * geometry.crop_patch_w - - if i != 0: - patch_idx[: geometry.left_margin, :] = -1 - if j != 0: - patch_idx[:, : geometry.left_margin] = -1 - if i != int(tiling[0]) - 1: - patch_idx[-geometry.right_margin :, :] = -1 - if j != int(tiling[1]) - 1: - patch_idx[:, -geometry.right_margin :] = -1 - patch_idx_arr[on_crop] = patch_idx - on_crop += 1 - return crop_arr, patch_idx_arr - - -def _build_overlapping_crops( - image: np.ndarray, - max_crops: int, - overlap_margins: list[int], - base_image_input_size: list[int], - image_mean: list[float], - image_std: list[float], - image_patch_size: int, -) -> tuple[np.ndarray, np.ndarray]: - geometry = _crop_geometry(overlap_margins, base_image_input_size, image_patch_size) - original_image_h, original_image_w = image.shape[:2] - - tiling = _select_tiling( - original_image_h - geometry.total_margin_pixels, - original_image_w - geometry.total_margin_pixels, - geometry.crop_window_size, - max_crops, - ) - - src = _resize_image( - image, - [ - tiling[0] * geometry.crop_window_size + geometry.total_margin_pixels, - tiling[1] * geometry.crop_window_size + geometry.total_margin_pixels, - ], - ) - src = _normalize_image(src, image_mean, image_std) - - crop_arr, patch_idx_arr = _fill_overlapping_crops(src, tiling, geometry) - patch_idx_arr = patch_idx_arr.reshape(int(tiling[0]), int(tiling[1]), geometry.crop_patch_h, geometry.crop_patch_w) - patch_idx_arr = patch_idx_arr.transpose(0, 2, 1, 3).reshape(-1) - patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape( - src.shape[0] // image_patch_size, - src.shape[1] // image_patch_size, - ) - return crop_arr, patch_idx_arr - - -def _batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray: - n_crops, h, w, c = array.shape - h_patches = h // patch_size - w_patches = w // patch_size - array = array.reshape(n_crops, h_patches, patch_size, w_patches, patch_size, c) - array = array.transpose(0, 1, 3, 2, 4, 5) - return array.reshape(n_crops, h_patches * w_patches, patch_size * patch_size * c) - - -def _arange_for_pooling(idx_arr: np.ndarray, pool_h: int, pool_w: int) -> np.ndarray: - h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0] - w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1] - idx_arr = np.pad( - idx_arr, - [[h_pad // 2, (h_pad + 1) // 2], [w_pad // 2, (w_pad + 1) // 2]], - mode="constant", - constant_values=-1, - ) - blocks_h = idx_arr.shape[0] // pool_h - blocks_w = idx_arr.shape[1] // pool_w - idx_arr = idx_arr.reshape(blocks_h, pool_h, blocks_w, pool_w) - idx_arr = idx_arr.transpose(0, 2, 1, 3) - return idx_arr.reshape(blocks_h, blocks_w, pool_h * pool_w) - - -def _image_to_patches_and_grids( - image: np.ndarray, - max_crops: int, - overlap_margins: list[int], - base_image_input_size: list[int], - image_mean: list[float], - image_std: list[float], - image_patch_size: int, - image_pooling_w: int, - image_pooling_h: int, - crop_mode: str, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - crop_patch_w = base_image_input_size[1] // image_patch_size - crop_patch_h = base_image_input_size[0] // image_patch_size - - if crop_mode == "resize": - resized, resize_idx = _build_resized_image( - image, - base_image_input_size, - image_mean, - image_std, - image_patch_size, - ) - resize_idx = _arange_for_pooling(resize_idx, image_pooling_h, image_pooling_w) - resized_h, resized_w = resize_idx.shape[:2] - resize_idx = resize_idx.reshape(-1, image_pooling_h * image_pooling_w) - image_grid = [np.asarray([resized_h, resized_w, 0, 0])] - return np.stack(image_grid, 0), _batch_pixels_to_patches(resized, image_patch_size), resize_idx - - if crop_mode not in {"overlap-and-resize-c2", "overlap-and-resize"}: - msg = f"Unsupported MolmoAct2 image crop_mode {crop_mode!r}." - raise ValueError(msg) - - crop_arr, patch_idx_arr = _build_overlapping_crops( - image, - max_crops, - overlap_margins, - base_image_input_size, - image_mean, - image_std, - image_patch_size, - ) - pooling_idx = _arange_for_pooling(patch_idx_arr, image_pooling_h, image_pooling_w) - h, w = pooling_idx.shape[:2] - pooling_idx = pooling_idx.reshape(-1, image_pooling_h * image_pooling_w) - - resized, resize_idx = _build_resized_image( - image, - base_image_input_size, - image_mean, - image_std, - image_patch_size, - ) - crop_arr = np.concatenate([resized, crop_arr], axis=0) - - resize_idx = _arange_for_pooling(resize_idx, image_pooling_h, image_pooling_w) - resized_h, resized_w = resize_idx.shape[:2] - resize_idx = resize_idx.reshape(-1, image_pooling_h * image_pooling_w) - - pooling_idx = np.where(pooling_idx >= 0, pooling_idx + crop_patch_h * crop_patch_w, -1) - pooling_idx = np.concatenate([resize_idx, pooling_idx], axis=0) - image_grid = [np.asarray([resized_h, resized_w, h, w])] - return np.stack(image_grid, 0), _batch_pixels_to_patches(crop_arr, image_patch_size), pooling_idx - - -def _to_hwc_uint8(images_bchw: np.ndarray) -> list[np.ndarray]: - out: list[np.ndarray] = [] - for image in images_bchw: - img = image - if np.issubdtype(img.dtype, np.floating): - if float(np.max(img)) <= 1.0: - img *= 255.0 - img = np.clip(img, 0.0, 255.0).astype(np.uint8) - elif img.dtype != np.uint8: - img = np.clip(img, 0, 255).astype(np.uint8) - out.append(np.transpose(img, (1, 2, 0))) - return out +_IMAGE_NDIM = 4 +_NUM_CHANNELS = 3 class MolmoAct2ImageProcessor: - """NumPy image processor producing MolmoAct2 patch tensors and pooling metadata.""" + """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, - do_convert_rgb: bool = True, - max_crops: int = 8, - overlap_margins: list[int] | None = None, - crop_mode: str = "overlap-and-resize-c2", - patch_size: int = 14, - pooling_size: list[int] | None = None, ) -> None: - """Initialize image processing options. - - Args: - size: Output image height and width. - image_mean: Per-channel image normalization means. - image_std: Per-channel image normalization standard deviations. - do_convert_rgb: Whether to convert images to RGB. - max_crops: Maximum number of overlapping crops. - overlap_margins: Left and right overlap margins in patches. - crop_mode: Image crop strategy. - patch_size: Height and width of each square image patch. - pooling_size: Width and height of each patch-pooling window. - """ - self.size = size if size is not None else {"height": 378, "width": 378} - self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5] - self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5] - self.do_convert_rgb = do_convert_rgb - self.max_crops = int(max_crops) - self.overlap_margins = overlap_margins if overlap_margins is not None else [4, 4] + """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.pooling_size = pooling_size if pooling_size is not None else [2, 2] + 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_bchw: np.ndarray) -> dict[str, np.ndarray]: - """Convert BCHW images into patches and pooling metadata. + def __call__(self, images: np.ndarray) -> dict[str, np.ndarray]: + """Return patches, pooling indices, grids, and crop counts. Returns: - Model-ready pixel values, pooling indices, image grids, and crop counts. - """ - image_list = _to_hwc_uint8(images_bchw) - patch_batches: list[np.ndarray] = [] - pooling_batches: list[np.ndarray] = [] - grids: list[np.ndarray] = [] - image_num_crops: list[int] = [] - - base_image_input_size = [int(self.size["height"]), int(self.size["width"])] - pool_h, pool_w = int(self.pooling_size[0]), int(self.pooling_size[1]) - - for image in image_list: - image_grid, crops, pooled_idx = _image_to_patches_and_grids( - image, - self.max_crops, - self.overlap_margins, - base_image_input_size, - self.image_mean, - self.image_std, - self.patch_size, - pool_w, - pool_h, - self.crop_mode, - ) - patch_batches.append(crops) - pooling_batches.append(pooled_idx) - grids.append(image_grid) - image_num_crops.append(int(crops.shape[0])) - - pixel_values = np.concatenate(patch_batches, axis=0) if patch_batches else np.zeros((0, 0, 0), dtype=np.float32) - image_token_pooling = ( - np.concatenate(pooling_batches, axis=0) - if pooling_batches - else np.zeros((0, pool_h * pool_w), dtype=np.int64) - ) - image_grids = np.concatenate(grids, axis=0) if grids else np.zeros((0, 4), dtype=np.int64) - image_num_crops_arr = np.asarray(image_num_crops, dtype=np.int64) + 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.astype(np.float32), - "image_token_pooling": image_token_pooling.astype(np.int64), - "image_grids": image_grids.astype(np.int64), - "image_num_crops": image_num_crops_arr, + "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 index 258c3348..b78f8c27 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -12,9 +12,18 @@ from __future__ import annotations -from dataclasses import dataclass, field +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 .molmoact2_image import MolmoAct2ImageProcessor + +_PACKED_IMAGE_NDIM = 5 @dataclass @@ -31,11 +40,11 @@ class MolmoAct2InputConfig: 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 = False + 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] = field(default_factory=list, init=False, repr=False) + image_token_ids: list[int] | None = None def __post_init__(self) -> None: """Collect the configured image token identifiers.""" @@ -49,16 +58,8 @@ def __post_init__(self) -> None: self.frame_end_token_id, self.image_low_res_id, ] - self._image_token_ids = [int(token_id) for token_id in ids if token_id is not None] - - @property - def image_token_ids(self) -> list[int]: - """Token ids that mark image content (for token type ids). - - Returns: - Configured image token identifiers. - """ - return self._image_token_ids + 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]: @@ -149,6 +150,7 @@ def expand_image_placeholders( 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) @@ -164,8 +166,9 @@ def expand_image_placeholders( else: expanded.append(token_int) expanded_rows.append(expanded) + expanded_widths.append(len(expanded) + int((~valid).sum())) - max_len = max((len(row) for row in expanded_rows), default=1) + max_len = max(expanded_widths, default=1) out_ids = np.full((len(expanded_rows), max_len), 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): @@ -200,7 +203,7 @@ def _batch_layout( _, 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), counts) + 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) @@ -323,8 +326,129 @@ def default_action_dim_is_pad(config: MolmoAct2InputConfig, *, batch_size: int) 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 + 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( + 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, + 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/tests/unit/inference/postprocessors/test_molmoact2.py b/tests/unit/inference/postprocessors/test_molmoact2.py index 972f223d..8241e874 100644 --- a/tests/unit/inference/postprocessors/test_molmoact2.py +++ b/tests/unit/inference/postprocessors/test_molmoact2.py @@ -35,6 +35,17 @@ def test_identity_without_stats(self) -> None: 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_uses_fixed_so101_transform_by_default(self) -> None: + checkpoint_values = [1.0, 88.0, 93.0] + processor = MolmoAct2Postprocessor( + action_stats={"q01": checkpoint_values, "q99": checkpoint_values}, + adapt_to_so101=True, + ) + + result = processor({ACTION: np.zeros((1, 3), dtype=np.float32)}) + + np.testing.assert_array_equal(result[ACTION], [[1.0, 2.0, 3.0]]) + def test_missing_action_raises(self) -> None: with pytest.raises(ValueError, match="action tensor"): MolmoAct2Postprocessor()({"other": np.zeros(1)}) diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 79df55ee..97c70c44 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -10,6 +10,7 @@ from physicalai.inference.manifest import ComponentSpec from physicalai.inference.component_factory import instantiate_component from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.preprocessors.molmoact2_image import MolmoAct2ImageProcessor from physicalai.inference.postprocessors import MolmoAct2Postprocessor @@ -94,6 +95,22 @@ def test_applies_masked_normalization_and_joint_transform(self) -> None: }) 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), @@ -113,6 +130,15 @@ def test_rejects_placeholder_image_mismatch(self) -> None: 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)}) @@ -121,6 +147,22 @@ 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 From 5bd0ac4593795907a2fe7b5e10b0357765f6aa56 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:49:08 +0200 Subject: [PATCH 16/22] fix: copilot comments --- pyproject.toml | 2 +- .../inference/preprocessors/molmoact2.py | 3 +- .../preprocessors/molmoact2_inputs.py | 24 +++++--- .../inference/preprocessors/test_molmoact2.py | 60 +++++++++++++++++++ .../preprocessors/test_stats_normalizer.py | 6 +- 5 files changed, 83 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 67bf1519..7140972d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] dependencies = [ - "numpy>=1.24", + "numpy>=1.24", "loguru>=0.7", # target PyTurboJPEG 1.x for libjpeg-turbo 2.x compatibility "PyTurboJPEG<2; sys_platform == 'linux'", diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2.py index 75240702..2c6bef78 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -252,7 +252,8 @@ def _raw_images(self, inputs: dict[str, Any]) -> list[np.ndarray]: if not self.image_keys and isinstance(container, np.ndarray): images = [container] elif not self.image_keys and isinstance(container, dict): - images = [np.asarray(value) for key, value in container.items() if "is_pad" not in str(key)] + 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] diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py index b78f8c27..7e474478 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -30,6 +30,7 @@ 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 @@ -145,8 +146,6 @@ def expand_image_placeholders( if int(image_grids.shape[0]) == 0: return input_ids, attention_mask, _build_token_type_ids(config, input_ids, attention_mask) - pad_values = input_ids[attention_mask == 0] - pad_token_id = int(pad_values[0]) if pad_values.size > 0 else 0 placeholder_id = int(config.image_placeholder_token_id) expanded_rows: list[list[int]] = [] @@ -169,7 +168,7 @@ def expand_image_placeholders( expanded_widths.append(len(expanded) + int((~valid).sum())) max_len = max(expanded_widths, default=1) - out_ids = np.full((len(expanded_rows), max_len), pad_token_id, dtype=input_ids.dtype) + 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: @@ -279,12 +278,22 @@ def build_batched_images( and ``(N, max_pooled, pool_area)``. Raises: - ValueError: If image-end token and image-grid counts differ. + ValueError: If image counts cannot be inferred from image-end tokens. """ - counts = (input_ids == int(config.image_end_token_id)).sum(1) + raw_counts = (input_ids == int(config.image_end_token_id)).sum(1) num_images = int(image_grids.shape[0]) - if int(counts.sum()) != num_images: - msg = f"image_end tokens ({int(counts.sum())}) do not match image grids ({num_images})." + 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) @@ -354,6 +363,7 @@ class MolmoAct2ModelInputs(Preprocessor): 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, diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 97c70c44..ef184a4a 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -10,6 +10,11 @@ from physicalai.inference.manifest import ComponentSpec from physicalai.inference.component_factory import instantiate_component from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.preprocessors.molmoact2_inputs import ( + MolmoAct2InputConfig, + build_batched_images, + expand_image_placeholders, +) from physicalai.inference.preprocessors.molmoact2_image import MolmoAct2ImageProcessor from physicalai.inference.postprocessors import MolmoAct2Postprocessor @@ -68,6 +73,21 @@ def test_builds_prompt_and_packs_ordered_cameras(self) -> None: 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) @@ -122,6 +142,46 @@ def test_preserves_masked_tokenizer_padding(self) -> None: 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()({ diff --git a/tests/unit/inference/preprocessors/test_stats_normalizer.py b/tests/unit/inference/preprocessors/test_stats_normalizer.py index 7986489a..4db666d8 100644 --- a/tests/unit/inference/preprocessors/test_stats_normalizer.py +++ b/tests/unit/inference/preprocessors/test_stats_normalizer.py @@ -173,8 +173,8 @@ def test_masked_dimensions_only(self) -> None: mode="quantiles", stats={ "state": { - "q01": [-1.0, -1.0, -1.0], - "q99": [1.0, 1.0, 1.0], + "q01": [0.0, 0.0, 0.0], + "q99": [2.0, 2.0, 2.0], "mask": [True, False, True], }, }, @@ -183,7 +183,7 @@ def test_masked_dimensions_only(self) -> None: result = normalizer(inputs) - np.testing.assert_allclose(result["state"], [[0.5, 4.0, -0.5]]) + np.testing.assert_allclose(result["state"], [[-0.5, 4.0, -1.5]]) class TestStatsNormalizerIdentity: From 1b9c5debe3b4678aa2dd0d5a6920c9544813e57e Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:51:29 +0200 Subject: [PATCH 17/22] fix: add token start and end for ov export --- src/physicalai/inference/preprocessors/molmoact2_inputs.py | 6 ++++++ tests/unit/inference/preprocessors/test_molmoact2.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2_inputs.py index 7e474478..9a1f2300 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2_inputs.py @@ -349,6 +349,9 @@ class MolmoAct2ModelInputs(Preprocessor): 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) @@ -370,6 +373,9 @@ def __post_init__(self) -> None: 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, diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index ef184a4a..8b786287 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -264,6 +264,9 @@ def __call__(self, tasks, **kwargs): 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), @@ -285,6 +288,9 @@ def __call__(self, tasks, **kwargs): 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"} From 68bed07a8722890d0bcc05bb668a3a8ae67954f2 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:52:42 +0200 Subject: [PATCH 18/22] feat: joint frame transformation processor, factor out for molmo --- src/physicalai/inference/component_factory.py | 2 + src/physicalai/inference/joint_transform.py | 59 +++++++++++++++++ .../inference/postprocessors/__init__.py | 2 + .../inference/postprocessors/joint_frame.py | 46 ++++++++++++++ .../inference/postprocessors/molmoact2.py | 10 +-- .../inference/preprocessors/__init__.py | 2 + .../inference/preprocessors/joint_frame.py | 50 +++++++++++++++ .../inference/preprocessors/molmoact2.py | 50 --------------- .../postprocessors/test_joint_frame.py | 37 +++++++++++ .../postprocessors/test_molmoact2.py | 25 +++----- .../preprocessors/test_joint_frame.py | 63 +++++++++++++++++++ .../inference/preprocessors/test_molmoact2.py | 19 +++--- 12 files changed, 283 insertions(+), 82 deletions(-) create mode 100644 src/physicalai/inference/joint_transform.py create mode 100644 src/physicalai/inference/postprocessors/joint_frame.py create mode 100644 src/physicalai/inference/preprocessors/joint_frame.py create mode 100644 tests/unit/inference/postprocessors/test_joint_frame.py create mode 100644 tests/unit/inference/preprocessors/test_joint_frame.py diff --git a/src/physicalai/inference/component_factory.py b/src/physicalai/inference/component_factory.py index 07be3428..ffef2676 100644 --- a/src/physicalai/inference/component_factory.py +++ b/src/physicalai/inference/component_factory.py @@ -114,11 +114,13 @@ def __repr__(self) -> str: component_registry.register("pi05", "physicalai.inference.preprocessors.Pi05Preprocessor") 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") def resolve_artifact(spec: ComponentSpec, export_dir: Path) -> ComponentSpec: diff --git a/src/physicalai/inference/joint_transform.py b/src/physicalai/inference/joint_transform.py new file mode 100644 index 00000000..159f5890 --- /dev/null +++ b/src/physicalai/inference/joint_transform.py @@ -0,0 +1,59 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Joint calibration frame transforms.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Sequence + + +class JointFrameTransform: + """Map leading joint values between robot and checkpoint calibration frames.""" + + 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 to_checkpoint(self, values: np.ndarray) -> np.ndarray: + """Map robot-frame joints to the checkpoint frame. + + Returns: + A transformed copy of ``values``. + """ + return self._apply(values, inverse=False) + + def to_robot(self, values: np.ndarray) -> np.ndarray: + """Map checkpoint-frame joints back to the robot frame. + + Returns: + A 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 diff --git a/src/physicalai/inference/postprocessors/__init__.py b/src/physicalai/inference/postprocessors/__init__.py index e4e52438..f38feab6 100644 --- a/src/physicalai/inference/postprocessors/__init__.py +++ b/src/physicalai/inference/postprocessors/__init__.py @@ -9,12 +9,14 @@ 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..a91b50ad --- /dev/null +++ b/src/physicalai/inference/postprocessors/joint_frame.py @@ -0,0 +1,46 @@ +# 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.joint_transform import JointFrameTransform +from physicalai.inference.postprocessors.base import Postprocessor + +if TYPE_CHECKING: + from collections.abc import Sequence + + +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.to_robot(np.asarray(outputs[self._feature])) + return result + + +__all__ = ["JointFramePostprocessor"] diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py index c5bfa018..a96cdd6c 100644 --- a/src/physicalai/inference/postprocessors/molmoact2.py +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -13,20 +13,17 @@ 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 JointFrameTransform, normalization_stats +from physicalai.inference.preprocessors.molmoact2 import normalization_stats class MolmoAct2Postprocessor(Postprocessor): - """Clamp, denormalize, and optionally convert actions to robot frame.""" + """Clamp and denormalize MolmoAct2 actions.""" def __init__( self, *, action_stats: dict[str, Any] | None = None, normalization_mode: str = "QUANTILES", - adapt_to_so101: bool = False, - joint_signs: list[float] | None = None, - joint_offsets: list[float] | None = None, ) -> None: """Store action postprocessing settings.""" self.denormalizer = ( @@ -38,7 +35,6 @@ def __init__( if action_stats else None ) - self.joint_transform = JointFrameTransform(joint_signs, joint_offsets) if adapt_to_so101 else None @override def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: @@ -58,8 +54,6 @@ def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: action = np.clip(np.asarray(action), -1.0, 1.0) if self.denormalizer is not None: action = self.denormalizer({ACTION: action})[ACTION] - if self.joint_transform is not None: - action = self.joint_transform.apply(action, inverse=True) result.pop("actions", None) result[ACTION] = action return result diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index 98134715..270a1632 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -9,6 +9,7 @@ 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 MolmoAct2Preprocessor from physicalai.inference.preprocessors.molmoact2_inputs import MolmoAct2ModelInputs @@ -22,6 +23,7 @@ __all__ = [ "HFTokenizer", + "JointFramePreprocessor", "LambdaPreprocessor", "MolmoAct2ModelInputs", "MolmoAct2Preprocessor", diff --git a/src/physicalai/inference/preprocessors/joint_frame.py b/src/physicalai/inference/preprocessors/joint_frame.py new file mode 100644 index 00000000..857da372 --- /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.joint_transform 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.to_checkpoint(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.py b/src/physicalai/inference/preprocessors/molmoact2.py index 2c6bef78..30cb1e6d 100644 --- a/src/physicalai/inference/preprocessors/molmoact2.py +++ b/src/physicalai/inference/preprocessors/molmoact2.py @@ -16,9 +16,6 @@ from physicalai.inference.preprocessors.base import Preprocessor from physicalai.inference.preprocessors.stats_normalizer import StatsNormalizer -SO101_JOINT_SIGNS = (1.0, -1.0, 1.0, 1.0, 1.0, 1.0) -SO101_JOINT_OFFSETS = (0.0, 90.0, 90.0, 0.0, 0.0, 0.0) - _TRAILING_PUNCTUATION = ".,!?;:" _IMAGE_NDIM = 4 _RGB_CHANNELS = 3 @@ -93,44 +90,6 @@ def _build_prompt( return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" -class JointFrameTransform: - """Map leading joints between robot and checkpoint frames.""" - - def __init__( - self, - signs: list[float] | None = None, - offsets: list[float] | None = None, - ) -> None: - """Use Studio SO101 defaults unless compatible overrides are supplied. - - Raises: - ValueError: If signs and offsets have different lengths. - """ - signs = list(SO101_JOINT_SIGNS) if signs is None else signs - offsets = list(SO101_JOINT_OFFSETS) if offsets is None else offsets - if len(signs) != len(offsets): - msg = f"joint_signs ({len(signs)}) and joint_offsets ({len(offsets)}) must match" - raise ValueError(msg) - self.signs = np.asarray(signs, dtype=np.float32) - self.offsets = np.asarray(offsets, dtype=np.float32) - - def apply(self, values: np.ndarray, *, inverse: bool) -> np.ndarray: - """Apply the forward or inverse affine transform. - - Returns: - A transformed copy of the input values. - """ - 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 MolmoAct2Preprocessor(Preprocessor): """Prepare normalized prompts and packed images before tokenization.""" @@ -146,9 +105,6 @@ def __init__( control_mode: str = "", add_setup_tokens: bool = True, add_control_tokens: bool = True, - adapt_to_so101: bool = False, - joint_signs: list[float] | None = None, - joint_offsets: list[float] | None = None, ) -> None: """Store observation preprocessing settings. @@ -165,7 +121,6 @@ def __init__( self.control_mode = control_mode self.add_setup_tokens = add_setup_tokens self.add_control_tokens = add_control_tokens - self.joint_transform = JointFrameTransform(joint_signs, joint_offsets) if adapt_to_so101 else None self.normalizer = ( StatsNormalizer( stats={STATE: normalization_stats(state_stats)}, @@ -194,8 +149,6 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: state = np.asarray(state, dtype=np.float32) if state.ndim == 1: state = state[None, :] - if self.joint_transform is not None: - state = self.joint_transform.apply(state, inverse=False) if self.normalizer is not None: state = self.normalizer({STATE: state})[STATE] state = np.clip(state, -1.0, 1.0) @@ -296,9 +249,6 @@ def _resize(self, images: np.ndarray) -> np.ndarray: __all__ = [ - "SO101_JOINT_OFFSETS", - "SO101_JOINT_SIGNS", - "JointFrameTransform", "MolmoAct2Preprocessor", "normalization_stats", ] 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..92ce7a04 --- /dev/null +++ b/tests/unit/inference/postprocessors/test_joint_frame.py @@ -0,0 +1,37 @@ +# 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 + + +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 index 8241e874..deabdc09 100644 --- a/tests/unit/inference/postprocessors/test_molmoact2.py +++ b/tests/unit/inference/postprocessors/test_molmoact2.py @@ -9,23 +9,25 @@ 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 MolmoAct2Postprocessor +from physicalai.inference.postprocessors import JointFramePostprocessor, MolmoAct2Postprocessor class TestMolmoAct2Postprocessor: - def test_clamps_masked_denormalizes_and_transforms(self) -> None: + def test_clamps_and_masked_denormalizes_before_joint_transform(self) -> None: processor = MolmoAct2Postprocessor( action_stats={ "q01": [0.0, 0.0, 0.0], "q99": [2.0, 2.0, 2.0], "mask": [True, False, True], }, - adapt_to_so101=True, - joint_signs=[1.0, -1.0], - joint_offsets=[0.0, 2.0], + ) + joint_transform = JointFramePostprocessor( + feature=ACTION, + signs=[1.0, -1.0], + offsets=[0.0, 2.0], ) - result = processor({"actions": np.array([[[2.0, 0.5, -2.0]]], dtype=np.float32)}) + 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 @@ -35,17 +37,6 @@ def test_identity_without_stats(self) -> None: 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_uses_fixed_so101_transform_by_default(self) -> None: - checkpoint_values = [1.0, 88.0, 93.0] - processor = MolmoAct2Postprocessor( - action_stats={"q01": checkpoint_values, "q99": checkpoint_values}, - adapt_to_so101=True, - ) - - result = processor({ACTION: np.zeros((1, 3), dtype=np.float32)}) - - np.testing.assert_array_equal(result[ACTION], [[1.0, 2.0, 3.0]]) - def test_missing_action_raises(self) -> None: with pytest.raises(ValueError, match="action tensor"): MolmoAct2Postprocessor()({"other": np.zeros(1)}) 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..a776e36f --- /dev/null +++ b/tests/unit/inference/preprocessors/test_joint_frame.py @@ -0,0 +1,63 @@ +# 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.joint_transform import JointFrameTransform +from physicalai.inference.manifest import ComponentSpec +from physicalai.inference.preprocessors import JointFramePreprocessor + + +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.to_checkpoint(robot_values) + + np.testing.assert_array_equal(checkpoint_values, [[12.0, 17.0, 4.0]]) + np.testing.assert_array_equal(transform.to_robot(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_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 index 8b786287..89b9602e 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -9,7 +9,7 @@ from physicalai.inference.constants import IMAGES, STATE, TASK, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK from physicalai.inference.manifest import ComponentSpec from physicalai.inference.component_factory import instantiate_component -from physicalai.inference.preprocessors import MolmoAct2ModelInputs, MolmoAct2Preprocessor +from physicalai.inference.preprocessors import JointFramePreprocessor, MolmoAct2ModelInputs, MolmoAct2Preprocessor from physicalai.inference.preprocessors.molmoact2_inputs import ( MolmoAct2InputConfig, build_batched_images, @@ -99,20 +99,25 @@ def test_accepts_batched_channels_last_camera_frames(self) -> None: assert float(result[IMAGES][0].max()) == 0.0 assert float(result[IMAGES][1].min()) == 1.0 - def test_applies_masked_normalization_and_joint_transform(self) -> None: + 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]}, - adapt_to_so101=True, - joint_signs=[1.0, -1.0], - joint_offsets=[0.0, 2.0], ) - result = processor({ + 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: From 868ea1129d5620b78dcffd3c31d0aa0fd0eb4444 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:03:45 +0200 Subject: [PATCH 19/22] refactor: forward and inverse joint transformations --- src/physicalai/inference/joint_transform.py | 12 ++++++------ .../inference/postprocessors/joint_frame.py | 2 +- .../inference/preprocessors/joint_frame.py | 2 +- .../unit/inference/preprocessors/test_joint_frame.py | 4 ++-- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/physicalai/inference/joint_transform.py b/src/physicalai/inference/joint_transform.py index 159f5890..2719e124 100644 --- a/src/physicalai/inference/joint_transform.py +++ b/src/physicalai/inference/joint_transform.py @@ -14,7 +14,7 @@ class JointFrameTransform: - """Map leading joint values between robot and checkpoint calibration frames.""" + """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. @@ -31,19 +31,19 @@ def __init__(self, *, signs: Sequence[float], offsets: Sequence[float]) -> None: self._signs = np.asarray(signs, dtype=np.float32) self._offsets = np.asarray(offsets, dtype=np.float32) - def to_checkpoint(self, values: np.ndarray) -> np.ndarray: - """Map robot-frame joints to the checkpoint frame. + 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 to_robot(self, values: np.ndarray) -> np.ndarray: - """Map checkpoint-frame joints back to the robot frame. + def inverse(self, values: np.ndarray) -> np.ndarray: + """Apply ``sign * (value - offset)`` to leading joint values. Returns: - A transformed copy of ``values``. + An inverse-transformed copy of ``values``. """ return self._apply(values, inverse=True) diff --git a/src/physicalai/inference/postprocessors/joint_frame.py b/src/physicalai/inference/postprocessors/joint_frame.py index a91b50ad..69c1931d 100644 --- a/src/physicalai/inference/postprocessors/joint_frame.py +++ b/src/physicalai/inference/postprocessors/joint_frame.py @@ -39,7 +39,7 @@ def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: msg = f"Joint frame postprocessor expected feature {self._feature!r}" raise ValueError(msg) result = dict(outputs) - result[self._feature] = self._transform.to_robot(np.asarray(outputs[self._feature])) + result[self._feature] = self._transform.inverse(np.asarray(outputs[self._feature])) return result diff --git a/src/physicalai/inference/preprocessors/joint_frame.py b/src/physicalai/inference/preprocessors/joint_frame.py index 857da372..9f14e892 100644 --- a/src/physicalai/inference/preprocessors/joint_frame.py +++ b/src/physicalai/inference/preprocessors/joint_frame.py @@ -34,7 +34,7 @@ def __call__(self, inputs: dict[str, Any]) -> dict[str, Any]: """ key = self._resolve_key(inputs) outputs = dict(inputs) - outputs[key] = self._transform.to_checkpoint(np.asarray(inputs[key])) + outputs[key] = self._transform.forward(np.asarray(inputs[key])) return outputs def _resolve_key(self, inputs: dict[str, Any]) -> str: diff --git a/tests/unit/inference/preprocessors/test_joint_frame.py b/tests/unit/inference/preprocessors/test_joint_frame.py index a776e36f..1d2f3340 100644 --- a/tests/unit/inference/preprocessors/test_joint_frame.py +++ b/tests/unit/inference/preprocessors/test_joint_frame.py @@ -14,10 +14,10 @@ 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.to_checkpoint(robot_values) + 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.to_robot(checkpoint_values), robot_values) + np.testing.assert_array_equal(transform.inverse(checkpoint_values), robot_values) def test_joint_transform_rejects_invalid_frame() -> None: From 628cf1c1411b63fe315ad9dbd16b89df34689bc7 Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:19:14 +0200 Subject: [PATCH 20/22] refactor: move joint_transform to postprocess --- src/physicalai/inference/joint_transform.py | 59 ------------------- .../inference/postprocessors/joint_frame.py | 47 ++++++++++++++- .../inference/preprocessors/joint_frame.py | 2 +- .../postprocessors/test_joint_frame.py | 18 ++++++ .../preprocessors/test_joint_frame.py | 18 ------ 5 files changed, 65 insertions(+), 79 deletions(-) delete mode 100644 src/physicalai/inference/joint_transform.py diff --git a/src/physicalai/inference/joint_transform.py b/src/physicalai/inference/joint_transform.py deleted file mode 100644 index 2719e124..00000000 --- a/src/physicalai/inference/joint_transform.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (C) 2026 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -"""Joint calibration frame transforms.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import numpy as np - -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 diff --git a/src/physicalai/inference/postprocessors/joint_frame.py b/src/physicalai/inference/postprocessors/joint_frame.py index 69c1931d..55677410 100644 --- a/src/physicalai/inference/postprocessors/joint_frame.py +++ b/src/physicalai/inference/postprocessors/joint_frame.py @@ -10,13 +10,58 @@ import numpy as np from typing_extensions import override -from physicalai.inference.joint_transform import JointFrameTransform 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.""" diff --git a/src/physicalai/inference/preprocessors/joint_frame.py b/src/physicalai/inference/preprocessors/joint_frame.py index 9f14e892..08ae19a6 100644 --- a/src/physicalai/inference/preprocessors/joint_frame.py +++ b/src/physicalai/inference/preprocessors/joint_frame.py @@ -10,7 +10,7 @@ import numpy as np from typing_extensions import override -from physicalai.inference.joint_transform import JointFrameTransform +from physicalai.inference.postprocessors.joint_frame import JointFrameTransform from physicalai.inference.preprocessors.base import Preprocessor if TYPE_CHECKING: diff --git a/tests/unit/inference/postprocessors/test_joint_frame.py b/tests/unit/inference/postprocessors/test_joint_frame.py index 92ce7a04..c64f6185 100644 --- a/tests/unit/inference/postprocessors/test_joint_frame.py +++ b/tests/unit/inference/postprocessors/test_joint_frame.py @@ -7,6 +7,24 @@ 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: diff --git a/tests/unit/inference/preprocessors/test_joint_frame.py b/tests/unit/inference/preprocessors/test_joint_frame.py index 1d2f3340..13ccc702 100644 --- a/tests/unit/inference/preprocessors/test_joint_frame.py +++ b/tests/unit/inference/preprocessors/test_joint_frame.py @@ -5,28 +5,10 @@ import pytest from physicalai.inference.component_factory import instantiate_component -from physicalai.inference.joint_transform import JointFrameTransform from physicalai.inference.manifest import ComponentSpec from physicalai.inference.preprocessors import JointFramePreprocessor -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_preprocessor_transforms_configured_feature() -> None: processor = instantiate_component( ComponentSpec( From 5ab81540fc8d1ad80e0821acd4b1ee5008904c9c Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:30:42 +0200 Subject: [PATCH 21/22] feat: add action key to postprocessor --- .../inference/postprocessors/molmoact2.py | 19 ++++++++++----- .../postprocessors/test_molmoact2.py | 23 +++++++++++++++---- .../inference/preprocessors/test_molmoact2.py | 1 + 3 files changed, 32 insertions(+), 11 deletions(-) diff --git a/src/physicalai/inference/postprocessors/molmoact2.py b/src/physicalai/inference/postprocessors/molmoact2.py index a96cdd6c..67b6ce01 100644 --- a/src/physicalai/inference/postprocessors/molmoact2.py +++ b/src/physicalai/inference/postprocessors/molmoact2.py @@ -22,10 +22,18 @@ class MolmoAct2Postprocessor(Postprocessor): def __init__( self, *, + action_key: str, action_stats: dict[str, Any] | None = None, normalization_mode: str = "QUANTILES", ) -> None: - """Store action postprocessing settings.""" + """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)}, @@ -44,17 +52,16 @@ def __call__(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: Outputs with the canonical denormalized action. Raises: - ValueError: If no action output is present. + ValueError: If the configured action output is absent. """ result = dict(outputs) - action = result.get(ACTION, result.get("actions")) - if action is None: - msg = "MolmoAct2 postprocessor expected an action tensor" + 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.pop("actions", None) result[ACTION] = action return result diff --git a/tests/unit/inference/postprocessors/test_molmoact2.py b/tests/unit/inference/postprocessors/test_molmoact2.py index deabdc09..b92ada40 100644 --- a/tests/unit/inference/postprocessors/test_molmoact2.py +++ b/tests/unit/inference/postprocessors/test_molmoact2.py @@ -15,6 +15,7 @@ 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], @@ -33,14 +34,26 @@ def test_clamps_and_masked_denormalizes_before_joint_transform(self) -> None: assert "actions" not in result def test_identity_without_stats(self) -> None: - processor = MolmoAct2Postprocessor() + 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="action tensor"): - MolmoAct2Postprocessor()({"other": np.zeros(1)}) + 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"), + ) - def test_registry_alias_instantiates(self) -> None: - processor = instantiate_component(ComponentSpec(type="molmoact2_postprocess")) 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/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 89b9602e..7780ec84 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -305,6 +305,7 @@ def __call__(self, tasks, **kwargs): postprocessor = instantiate_component( ComponentSpec( type="molmoact2_postprocess", + action_key="action", action_stats={"q01": [0.0, 0.0], "q99": [2.0, 2.0]}, ), ) From bcac1322bd8059766247f6beb887ef8376a3499f Mon Sep 17 00:00:00 2001 From: Alfie <51797647+alfieroddan@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:54:55 +0200 Subject: [PATCH 22/22] refactor: move molmo processors to own subfolder --- .../inference/postprocessors/molmoact2/__init__.py | 8 ++++++++ .../{molmoact2.py => molmoact2/processor.py} | 0 src/physicalai/inference/preprocessors/__init__.py | 3 +-- .../inference/preprocessors/molmoact2/__init__.py | 13 +++++++++++++ .../{molmoact2_image.py => molmoact2/image.py} | 0 .../{molmoact2_inputs.py => molmoact2/inputs.py} | 2 +- .../{molmoact2.py => molmoact2/processor.py} | 0 .../unit/inference/preprocessors/test_molmoact2.py | 8 ++++---- 8 files changed, 27 insertions(+), 7 deletions(-) create mode 100644 src/physicalai/inference/postprocessors/molmoact2/__init__.py rename src/physicalai/inference/postprocessors/{molmoact2.py => molmoact2/processor.py} (100%) create mode 100644 src/physicalai/inference/preprocessors/molmoact2/__init__.py rename src/physicalai/inference/preprocessors/{molmoact2_image.py => molmoact2/image.py} (100%) rename src/physicalai/inference/preprocessors/{molmoact2_inputs.py => molmoact2/inputs.py} (99%) rename src/physicalai/inference/preprocessors/{molmoact2.py => molmoact2/processor.py} (100%) 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.py b/src/physicalai/inference/postprocessors/molmoact2/processor.py similarity index 100% rename from src/physicalai/inference/postprocessors/molmoact2.py rename to src/physicalai/inference/postprocessors/molmoact2/processor.py diff --git a/src/physicalai/inference/preprocessors/__init__.py b/src/physicalai/inference/preprocessors/__init__.py index a32c6b3c..0485de9d 100644 --- a/src/physicalai/inference/preprocessors/__init__.py +++ b/src/physicalai/inference/preprocessors/__init__.py @@ -11,8 +11,7 @@ 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 MolmoAct2Preprocessor -from physicalai.inference.preprocessors.molmoact2_inputs import MolmoAct2ModelInputs +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 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 similarity index 100% rename from src/physicalai/inference/preprocessors/molmoact2_image.py rename to src/physicalai/inference/preprocessors/molmoact2/image.py diff --git a/src/physicalai/inference/preprocessors/molmoact2_inputs.py b/src/physicalai/inference/preprocessors/molmoact2/inputs.py similarity index 99% rename from src/physicalai/inference/preprocessors/molmoact2_inputs.py rename to src/physicalai/inference/preprocessors/molmoact2/inputs.py index 9a1f2300..c2db090e 100644 --- a/src/physicalai/inference/preprocessors/molmoact2_inputs.py +++ b/src/physicalai/inference/preprocessors/molmoact2/inputs.py @@ -21,7 +21,7 @@ from physicalai.inference.constants import IMAGES, TOKENIZED_PROMPT, TOKENIZED_PROMPT_MASK from physicalai.inference.preprocessors.base import Preprocessor -from .molmoact2_image import MolmoAct2ImageProcessor +from .image import MolmoAct2ImageProcessor _PACKED_IMAGE_NDIM = 5 diff --git a/src/physicalai/inference/preprocessors/molmoact2.py b/src/physicalai/inference/preprocessors/molmoact2/processor.py similarity index 100% rename from src/physicalai/inference/preprocessors/molmoact2.py rename to src/physicalai/inference/preprocessors/molmoact2/processor.py diff --git a/tests/unit/inference/preprocessors/test_molmoact2.py b/tests/unit/inference/preprocessors/test_molmoact2.py index 7780ec84..30aa6d27 100644 --- a/tests/unit/inference/preprocessors/test_molmoact2.py +++ b/tests/unit/inference/preprocessors/test_molmoact2.py @@ -6,17 +6,17 @@ 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.component_factory import instantiate_component +from physicalai.inference.postprocessors import MolmoAct2Postprocessor from physicalai.inference.preprocessors import JointFramePreprocessor, MolmoAct2ModelInputs, MolmoAct2Preprocessor -from physicalai.inference.preprocessors.molmoact2_inputs import ( +from physicalai.inference.preprocessors.molmoact2.image import MolmoAct2ImageProcessor +from physicalai.inference.preprocessors.molmoact2.inputs import ( MolmoAct2InputConfig, build_batched_images, expand_image_placeholders, ) -from physicalai.inference.preprocessors.molmoact2_image import MolmoAct2ImageProcessor -from physicalai.inference.postprocessors import MolmoAct2Postprocessor def _prepare(**kwargs) -> MolmoAct2Preprocessor: