# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json import os import re import warnings from collections.abc import Callable from copy import deepcopy from pathlib import Path from typing import Any import numpy as np import torch import torchvision.transforms.v2 as transforms from PIL import Image from transformers import AutoProcessor, ProcessorMixin from transformers.feature_extraction_utils import BatchFeature from transformers.utils import cached_file from vllm.logger import init_logger from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable class LetterBoxTransform: """Pad image to square dimensions by adding black bars to the smaller side.""" def __call__(self, img: torch.Tensor) -> torch.Tensor: import math *leading_dims, c, h, w = img.shape if h == w: return img max_dim = max(h, w) pad_h, pad_w = max_dim - h, max_dim - w pad_top, pad_left = pad_h // 2, pad_w // 2 pad_bottom, pad_right = pad_h - pad_top, pad_w - pad_left if leading_dims: batch_size = math.prod(leading_dims) img_r = img.reshape(batch_size, c, h, w) padded = transforms.functional.pad(img_r, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) return padded.reshape(leading_dims + [c, max_dim, max_dim]) return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) def _smallest_edge_resize(image: np.ndarray, max_size: int) -> np.ndarray: """albumentations 1.4.18 ``SmallestMaxSize(max_size, interpolation=cv2.INTER_AREA)``. Scale off the shortest edge, skip entirely at scale 1.0, round each dimension independently; ``cv2.resize`` takes ``dsize`` as ``(width, height)``. """ import cv2 height, width = image.shape[:2] scale = max_size / float(min(height, width)) if scale == 1.0: return image new_height, new_width = (round(dim * scale) for dim in (height, width)) if (new_height, new_width) == (height, width): return image return cv2.resize(image, dsize=(new_width, new_height), interpolation=cv2.INTER_AREA) def _fractional_center_crop(image: np.ndarray, fraction: float) -> np.ndarray: """Isaac-GR00T ``FractionalCenterCrop``: ``int()`` truncation (not rounding), floor 1, centred.""" height, width = image.shape[:2] crop_height = max(1, int(height * fraction)) crop_width = max(1, int(width * fraction)) y_min = (height - crop_height) // 2 x_min = (width - crop_width) // 2 return image[y_min : y_min + crop_height, x_min : x_min + crop_width] def _letter_box_pad(image: np.ndarray) -> np.ndarray: """Isaac-GR00T ``LetterBoxPad``: zero-pad to square, centred, remainder on the far side.""" import cv2 height, width = image.shape[:2] if height == width: return image longest = max(height, width) pad_h, pad_w = longest - height, longest - width top, left = pad_h // 2, pad_w // 2 return cv2.copyMakeBorder(image, top, pad_h - top, left, pad_w - left, cv2.BORDER_CONSTANT, value=0) class AlbumentationsEvalImageTransform: """Eval transform of the Isaac-GR00T ``use_albumentations=True`` recipe, without albumentations. ``[LetterBoxPad] -> SmallestMaxSize(INTER_AREA) -> FractionalCenterCrop -> SmallestMaxSize(INTER_AREA)``, transcribed from the eval branch of Isaac-GR00T ``build_image_transformations_albumentations`` and albumentations 1.4.18, so a checkpoint trained with that recipe sees the same pixels at inference. Per frame: HWC uint8 RGB array in, ``(C, H, W)`` uint8 tensor out (the ``_get_vlm_inputs`` contract). """ def __init__(self, max_size: int, fraction: float, letter_box_transform: bool): self.max_size = int(max_size) self.fraction = float(fraction) self.letter_box_transform = letter_box_transform def __call__(self, image: np.ndarray) -> torch.Tensor: array = np.asarray(image) if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3: raise ValueError( "GR00T eval image transform expects an HWC uint8 RGB frame, " f"got shape={tuple(array.shape)} dtype={array.dtype}" ) if self.letter_box_transform: array = _letter_box_pad(array) array = _smallest_edge_resize(array, self.max_size) array = _fractional_center_crop(array, self.fraction) array = _smallest_edge_resize(array, self.max_size) # The crop is a non-contiguous view; the caller stacks frames right after this. return torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1) def _build_eval_image_transform( image_target_size: list[int], image_crop_size: list[int], *, shortest_image_edge: int | None = None, crop_fraction: float | None = None, use_albumentations: bool = False, letter_box_transform: bool = False, ) -> Callable[[np.ndarray], torch.Tensor]: """Deterministic eval/inference image transform, selected the way Isaac-GR00T's processor does. ``use_albumentations=True``: ``[letterbox] -> resize (shortest edge) -> fractional centre crop -> resize`` (:class:`AlbumentationsEvalImageTransform`). Otherwise the torchvision recipe ``[letterbox] -> resize -> centre crop -> resize``. Letterboxing is gated by ``letter_box_transform`` on both paths. """ if use_albumentations: if crop_fraction is None: if image_crop_size is None or image_target_size is None: raise ValueError("image_crop_size and image_target_size are required when crop_fraction is None") fraction = image_crop_size[0] / image_target_size[0] else: fraction = crop_fraction if shortest_image_edge is None: if image_target_size is None: raise ValueError("image_target_size is required when shortest_image_edge is None") max_size = image_target_size[0] else: max_size = shortest_image_edge return AlbumentationsEvalImageTransform(max_size, fraction, letter_box_transform) steps: list[Callable[..., torch.Tensor]] = [transforms.ToImage()] if letter_box_transform: steps.append(LetterBoxTransform()) steps.extend( [ transforms.Resize(size=image_target_size), transforms.CenterCrop(size=image_crop_size), transforms.Resize(size=image_target_size), ] ) return transforms.Compose(steps) logger = init_logger(__name__) # Suppress protobuf deprecation warnings warnings.filterwarnings("ignore", category=DeprecationWarning, module="google.protobuf") EMBODIMENT_TAG_TO_PROJECTOR_INDEX = { # Pretrain embodiment ids "oxe_droid_relative_eef_relative_joint": 24, "xdof_relative_eef_relative_joint": 27, "xdof_relative_eef_relative_joint_subtask": 27, "real_g1_relative_eef_relative_joints": 25, "real_r1_pro_sharpa_relative_eef": 26, "real_r1_pro_sharpa_relative_eef_human": 26, "real_r1_pro_sharpa_relative_eef_maxinsights": 26, "real_r1_pro_sharpa_relative_eef_mecka": 26, # Posttrain embodiment ids "unitree_g1_full_body_with_waist_height_nav_cmd": 25, "unitree_g1_sonic": 11, "simpler_env_google": 0, "simpler_env_widowx": 1, "libero_sim": 2, "new_embodiment": 10, } QWEN3_VL_2B_PROCESSOR = "Qwen/Qwen3-VL-2B-Instruct" def build_processor(model_name: str, transformers_loading_kwargs: dict) -> ProcessorMixin: try: from transformers import Qwen3VLProcessor except ImportError as exc: raise ImportError( "GR00T-N1.7 requires transformers>=4.57.1 for Qwen3VLProcessor " "(the repo's declared floor is 4.56.0). Please upgrade transformers." ) from exc if model_name == "nvidia/Cosmos-Reason2-2B": # Cosmos-Reason2-2B lacks a Qwen3VLProcessor; fall back to upstream artifacts. logger.warning_once( "Substituting processor from %s because %s does not ship one. " "If you fine-tune Cosmos-Reason2-2B's tokenizer/image processor, " "load the processor explicitly instead of relying on this fallback.", QWEN3_VL_2B_PROCESSOR, model_name, ) model_name = QWEN3_VL_2B_PROCESSOR return Qwen3VLProcessor.from_pretrained(model_name, **transformers_loading_kwargs) class Gr00tN1d7DataCollator: def __init__( self, model_name: str, model_type: str = "qwen", transformers_loading_kwargs: dict | None = None, ): if transformers_loading_kwargs is None: transformers_loading_kwargs = {} self.processor = build_processor(model_name, transformers_loading_kwargs) self.processor.tokenizer.padding_side = "left" self.model_type = model_type self.model_name = model_name def __call__(self, features: list[dict[str, Any]]) -> BatchFeature: batch = {} keys = list(set().union(*(elem.keys() for elem in features))) for key in keys: values = [elem[key] for elem in features if key in elem] if key == "vlm_content": text_list = [] image_inputs = [] for v in values: curr_text_list = [v["text"]] text_list += curr_text_list curr_image_inputs = v["images"] image_inputs += curr_image_inputs vlm_inputs = self.processor( text=text_list, images=image_inputs, return_tensors="pt", padding=True, ) for k, v in vlm_inputs.items(): batch[k] = v elif key in ( "pixel_values", "image_grid_thw", "attention_mask", "input_ids", ): raise Exception("Not implemented") else: batch[key] = torch.from_numpy(np.stack(values)) return BatchFeature(data={"inputs": batch}) def __str__(self): return f"Gr00tN1d7DataCollator(model_name={self.model_name}, model_type={self.model_type})" class Gr00tN1d7Processor(ProcessorMixin): data_collator_class = Gr00tN1d7DataCollator def __init__( self, modality_configs: dict[str, dict[str, ModalityConfig]], statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None, use_percentiles: bool = False, clip_outliers: bool = True, image_crop_size: list[int] = None, image_target_size: list[int] = None, shortest_image_edge: int = 256, crop_fraction: float = 0.95, random_rotation_angle: int | None = None, color_jitter_params: dict[str, float] | None = None, formalize_language: bool = True, model_name: str = "nvidia/Cosmos-Reason2-2B", model_type: str = "qwen", max_state_dim: int = 29, max_action_dim: int = 29, max_action_horizon: int = 50, apply_sincos_state_encoding: bool = False, use_relative_action: bool = False, embodiment_id_mapping: dict[str, int] | None = None, transformers_loading_kwargs: dict | None = None, exclude_state: bool = False, # Normalization use_mean_std: bool = False, # Image pipeline selection; must match the recipe the checkpoint was trained with use_albumentations: bool = False, letter_box_transform: bool = False, **kwargs, # absorb deprecated training-only keys from saved processor_config.json ): if transformers_loading_kwargs is None: transformers_loading_kwargs = {"trust_remote_code": True} if kwargs: logger.debug("Gr00tN1d7Processor: ignoring unknown keys: %s", list(kwargs)) self.modality_configs = parse_modality_configs(modality_configs) # Initialize StateActionProcessor for state/action normalization self.state_action_processor = StateActionProcessor( modality_configs=modality_configs, statistics=statistics, use_percentiles=use_percentiles, clip_outliers=clip_outliers, apply_sincos_state_encoding=apply_sincos_state_encoding, use_relative_action=use_relative_action, ) self.use_percentiles = use_percentiles self.use_mean_std = use_mean_std self.clip_outliers = clip_outliers self.apply_sincos_state_encoding = apply_sincos_state_encoding self.use_relative_action = use_relative_action self.exclude_state = exclude_state self.formalize_language = formalize_language self.model_name = model_name self.model_type = model_type self.max_state_dim = max_state_dim self.max_action_dim = max_action_dim self.max_action_horizon = max_action_horizon self.image_crop_size = image_crop_size self.image_target_size = image_target_size self.random_rotation_angle = random_rotation_angle self.color_jitter_params = color_jitter_params self.processor = build_processor(model_name, transformers_loading_kwargs) self.processor.tokenizer.padding_side = "left" self.embodiment_id_mapping = embodiment_id_mapping or EMBODIMENT_TAG_TO_PROJECTOR_INDEX for k, v in EMBODIMENT_TAG_TO_PROJECTOR_INDEX.items(): if k not in self.embodiment_id_mapping: self.embodiment_id_mapping[k] = v self.shortest_image_edge = shortest_image_edge self.crop_fraction = crop_fraction self.use_albumentations = use_albumentations self.letter_box_transform = letter_box_transform self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} # Eval/inference image transform self.eval_image_transform = _build_eval_image_transform( image_target_size, image_crop_size, shortest_image_edge=shortest_image_edge, crop_fraction=crop_fraction, use_albumentations=use_albumentations, letter_box_transform=letter_box_transform, ) self._collator = self.data_collator_class( model_name=model_name, model_type=model_type, transformers_loading_kwargs=transformers_loading_kwargs, ) @property def collator(self): return self._collator def set_statistics( self, statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]], override: bool = False, ) -> None: """Set dataset statistics for normalization.""" for key in statistics: if key not in self.statistics or override: if override: logger.info("Overriding statistics for %s", key) self.statistics[key] = deepcopy(statistics[key]) else: logger.debug("Embodiment tag %s already in statistics, skipping update", key) self.state_action_processor.set_statistics(statistics, override=override) # Compute action dimensions for convenience self.action_dim = {} for embodiment_tag in self.state_action_processor.statistics: self.action_dim[embodiment_tag] = self.state_action_processor.get_action_dim(embodiment_tag) def decode_action( self, action: np.ndarray, embodiment_tag: EmbodimentTag, state: dict[str, np.ndarray] | None = None, ): """Undo action normalization and convert relative actions to absolute.""" # Split concatenated action into joint groups out_dict = {} start_idx = 0 joint_groups = self.modality_configs[embodiment_tag.value]["action"].modality_keys action_horizon = len(self.modality_configs[embodiment_tag.value]["action"].delta_indices) for key in joint_groups: joint_dim = self.state_action_processor.norm_params[embodiment_tag.value]["action"][key]["dim"].item() out_dict[key] = action[..., :action_horizon, start_idx : start_idx + joint_dim] start_idx += joint_dim # Use StateActionProcessor to unnormalize and convert to absolute return self.state_action_processor.unapply_action(out_dict, embodiment_tag.value, state=state) def _apply_vlm_processing(self, images: np.ndarray, language: str) -> BatchFeature: pil_images = [Image.fromarray(np.transpose(v, (1, 2, 0))) for v in images] conversation = [ { "role": "user", "content": [ *[{"type": "image", "image": img} for img in pil_images], {"type": "text", "text": language}, ], } ] text = self.processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False) return { "vlm_content": { "text": text, "images": pil_images, "conversation": conversation, } } def __call__( self, messages: list[dict[str, Any]], ): assert len(messages) == 1 content = messages[0]["content"] embodiment_tag = content.embodiment action_data = content.actions state_data = content.states norm_state_dict, normalized_actions = self.state_action_processor.apply( state=state_data, action=action_data, embodiment_tag=embodiment_tag.value, ) if normalized_actions: action_keys = self.modality_configs[embodiment_tag.value]["action"].modality_keys normalized_actions = torch.cat( [torch.from_numpy(normalized_actions[key]) for key in action_keys], dim=-1, ) # (t, d) action_dim = normalized_actions.shape[1] normalized_actions = torch.cat( [ normalized_actions, torch.zeros( normalized_actions.shape[0], self.max_action_dim - normalized_actions.shape[1], ), ], dim=-1, ) # (t, max_action_dim) action_horizon = normalized_actions.shape[0] assert action_horizon <= self.max_action_horizon, ( f"Action sequence length {action_horizon} exceeds max_action_horizon" f" {self.max_action_horizon}. Increase model config action_horizon to" f" >= {action_horizon}." ) normalized_actions = torch.cat( [ normalized_actions, torch.zeros( self.max_action_horizon - normalized_actions.shape[0], self.max_action_dim, ), ], dim=0, ) # (max_action_horizon, max_action_dim) action_mask = torch.ones_like(normalized_actions) action_mask[action_horizon:] = 0 action_mask[:, action_dim:] = 0 else: normalized_actions = None action_mask = None state_keys = self.modality_configs[embodiment_tag.value]["state"].modality_keys exclude_state = self.exclude_state or getattr( self.modality_configs[embodiment_tag.value]["state"], "exclude_state", False ) if exclude_state: normalized_states = torch.cat( [torch.from_numpy(np.zeros_like(state_data[key])) for key in state_keys], dim=-1 ) else: normalized_states = torch.cat([torch.from_numpy(norm_state_dict[key]) for key in state_keys], dim=-1) normalized_states = torch.cat( [ normalized_states, torch.zeros( normalized_states.shape[0], self.max_state_dim - normalized_states.shape[1], ), ], dim=-1, ) image_transform = self.eval_image_transform image_keys = self.modality_configs[embodiment_tag.value]["video"].modality_keys if self.formalize_language: language = content.text.lower() language = re.sub(r"[^\w\s]", "", language) else: language = content.text vlm_inputs = self._get_vlm_inputs( image_keys=image_keys, images=content.images, masks=content.masks, image_transform=image_transform, language=language, ) transformed_inputs = { "state": normalized_states.to(torch.get_default_dtype()), } if normalized_actions is not None: transformed_inputs["action"] = normalized_actions.to(torch.get_default_dtype()) # Add VLM inputs transformed_inputs.update(vlm_inputs) if action_mask is not None: transformed_inputs["action_mask"] = action_mask transformed_inputs["embodiment_id"] = self.embodiment_id_mapping[embodiment_tag.value] return transformed_inputs def _get_vlm_inputs( self, image_keys: list[str], images: list[Image.Image], masks: dict[str, list[np.ndarray]] | None, image_transform: Callable[[np.ndarray], torch.Tensor], language: str, ): temporal_stacked_images = {} if masks is not None: raise ValueError("Mask-based transforms are not supported at inference.") for view in image_keys: assert view in images, f"{view} not in {images}" temporal_stacked_images[view] = torch.stack([image_transform(img) for img in images[view]]) # (T, C, H, W) for k, v in temporal_stacked_images.items(): assert isinstance(k, str), f"{k} is not a string" assert isinstance(v, torch.Tensor), f"{v} is not a torch tensor" assert v.ndim == 4, f"{v} is not a 4D tensor" assert v.dtype == torch.uint8, f"{v} is not a uint8 tensor" assert v.shape[1] == 3, f"{v} is not a 3 channel tensor" stacked_images = ( torch.stack([temporal_stacked_images[view] for view in image_keys], dim=1).flatten(0, 1).numpy() ) # (T*V, C, H, W), processor expects numpy array vlm_inputs = self._apply_vlm_processing(stacked_images, language) return vlm_inputs def save_pretrained(self, save_directory: str | Path) -> list[Path]: save_directory = Path(save_directory) save_directory.mkdir(parents=True, exist_ok=True) main_config_file = save_directory / "processor_config.json" statistics_file = save_directory / "statistics.json" embodiment_id_file = save_directory / "embodiment_id.json" config = { "processor_class": self.__class__.__name__, "processor_kwargs": { "modality_configs": to_json_serializable(self.modality_configs), "image_crop_size": self.image_crop_size, "image_target_size": self.image_target_size, "random_rotation_angle": self.random_rotation_angle, "color_jitter_params": self.color_jitter_params, "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, "use_albumentations": self.use_albumentations, "letter_box_transform": self.letter_box_transform, "model_name": self.model_name, "model_type": self.model_type, "formalize_language": self.formalize_language, "max_state_dim": self.max_state_dim, "max_action_dim": self.max_action_dim, "max_action_horizon": self.max_action_horizon, "use_percentiles": self.use_percentiles, "use_mean_std": self.use_mean_std, "clip_outliers": self.clip_outliers, "apply_sincos_state_encoding": self.apply_sincos_state_encoding, "use_relative_action": self.use_relative_action, "exclude_state": self.exclude_state, }, } with open(main_config_file, "w") as f: json.dump(config, f, indent=2) with open(statistics_file, "w") as f: json.dump( to_json_serializable(self.state_action_processor.statistics), f, indent=2, ) with open(embodiment_id_file, "w") as f: json.dump(self.embodiment_id_mapping, f, indent=2) return [main_config_file, statistics_file, embodiment_id_file] @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", {"trust_remote_code": True}) pretrained_model_name_or_path = Path(pretrained_model_name_or_path) config_file = pretrained_model_name_or_path / "processor_config.json" statistics_file = pretrained_model_name_or_path / "statistics.json" embodiment_id_file = pretrained_model_name_or_path / "embodiment_id.json" is_local = os.path.isdir(pretrained_model_name_or_path) if not is_local: config_file = Path(cached_file(pretrained_model_name_or_path, "processor_config.json")) statistics_file = Path(cached_file(pretrained_model_name_or_path, "statistics.json")) embodiment_id_file = Path(cached_file(pretrained_model_name_or_path, "embodiment_id.json")) with open(config_file) as f: config = json.load(f) with open(statistics_file) as f: statistics = json.load(f) if embodiment_id_file.exists(): with open(embodiment_id_file) as f: embodiment_id_mapping = json.load(f) else: embodiment_id_mapping = None processor_kwargs = config["processor_kwargs"] processor_kwargs["statistics"] = statistics processor_kwargs["embodiment_id_mapping"] = embodiment_id_mapping # Backfill missing fields from older checkpoints. processor_kwargs.setdefault("model_name", "nvidia/Cosmos-Reason2-2B") processor_kwargs.setdefault("model_type", "qwen") processor_kwargs.setdefault("clip_outliers", True) # Directly override other processor kwargs if kwargs: # Override modality configs while keeping pretrained embodiment configs modality_configs = kwargs.pop("modality_configs", {}) for embodiment_tag, modality_config in modality_configs.items(): processor_kwargs["modality_configs"][embodiment_tag] = modality_config override_keys = [ "random_rotation_angle", "color_jitter_params", "use_relative_action", "exclude_state", "use_mean_std", "model_name", "model_type", "max_action_horizon", "max_state_dim", "max_action_dim", ] for key in override_keys: if key in kwargs: override = kwargs.pop(key) if override is not None: processor_kwargs[key] = override return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs) # transformers >= 5.5 requires the config *class* (not a model_type string) as the # first arg to AutoProcessor.register — it does `key.__module__` internally. Passing # the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute # '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in # modeling/gr00t_n1d7.py. AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)