tzhouamandClaude Opus 4.8
b5c6a30d4e
Three regressions surfaced only on dev/vllm-align in Buildkite build 2520 (26 failing jobs vs 5 on main); ~15 of the 19 align-unique failures trace to these: 1. numpy 2.5 vs numba: Dockerfile.ci pinned numpy==2.2.6 before the flashinfer and torch cu130 reinstalls, which resolved numpy back to 2.5 and broke 'import numba' (needs numpy <= 2.4). Re-pin numpy as the final dependency op. 2. libnvJitLink.so.13 not found: the CUDA-13 JIT linker lives in the nvidia-cuda-nvjitlink-cu13 wheel, not /usr/local/cuda (CUDA 12) as the old LD_LIBRARY_PATH assumed; bitsandbytes/torch.compile subprocesses could not dlopen it (broke Entrypoint L4, a multi-GPU diffusion job, both Quantization jobs via hang/timeout). Install the wheel and register its lib dir via ldconfig. 3. transformers >= 5.5 AutoProcessor.register now requires the config class, not a model_type string; processing_gr00t_n1d7.py passed "Gr00tN1d7" -> AttributeError 'str' object has no attribute '__module__' -> pytest collection errors (exit 2). Pass Gr00tN1d7Config, matching AutoModel.register in modeling/gr00t_n1d7.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
562 lines
23 KiB
Python
562 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import warnings
|
|
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 _build_eval_image_transform(
|
|
image_target_size: list[int],
|
|
image_crop_size: list[int],
|
|
) -> transforms.Compose:
|
|
"""Deterministic eval/inference image transform (letterbox → resize → centercrop → resize)."""
|
|
return transforms.Compose(
|
|
[
|
|
transforms.ToImage(),
|
|
LetterBoxTransform(),
|
|
transforms.Resize(size=image_target_size),
|
|
transforms.CenterCrop(size=image_crop_size),
|
|
transforms.Resize(size=image_target_size),
|
|
]
|
|
)
|
|
|
|
|
|
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,
|
|
**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.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,
|
|
)
|
|
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: transforms.Compose,
|
|
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,
|
|
"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)
|