Zenma/VLA-Adapter-LIBERO-Spatial-5000
015
1"""2modeling_prismatic.py3 4Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.5Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,6but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.7"""8 9import logging10from dataclasses import dataclass11from functools import partial12from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union13import numpy as np14import timm15import tokenizers16import torch17import torch.nn as nn18import transformers19from timm.models.vision_transformer import LayerScale20from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel21from transformers.modeling_outputs import ModelOutput22 23from prismatic.training.train_utils import (24 get_current_action_mask,25 get_next_actions_mask,26)27from prismatic.vla.constants import (28 ACTION_DIM,29 ACTION_PROPRIO_NORMALIZATION_TYPE,30 ACTION_TOKEN_BEGIN_IDX,31 IGNORE_INDEX,32 NUM_ACTIONS_CHUNK,33 STOP_INDEX,34 NormalizationType,35 NUM_TOKENS36)37from .configuration_prismatic import OpenVLAConfig, PrismaticConfig38 39 40 41# Set up logger42logger = logging.getLogger(__name__)43 44 45# === Utility Functions for Monkey-Patching ===46def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:47 def wrapper(*args: Any, **kwargs: Any) -> Any:48 result = fn(*args, **kwargs)49 return result[0] if isinstance(result, tuple) else result50 51 return wrapper52 53 54 55# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.56# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L10957# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L396058def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:59 return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor60 61 62 63def ls_apply_patch(ls_module: LayerScale):64 ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())65 ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)66 del ls_module.gamma67 68 69 70# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===71class PrismaticVisionBackbone(nn.Module):72 """73 Vision backbone for Prismatic models that handles image feature extraction.74 75 Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.76 For fused backbones, features from both models are concatenated along the feature dimension.77 """78 79 def __init__(80 self,81 use_fused_vision_backbone: bool,82 image_sizes: List[int],83 timm_model_ids: List[str],84 timm_override_act_layers: List[Optional[str]],85 ) -> None:86 """87 Initialize the vision backbone.88 89 Args:90 use_fused_vision_backbone: Whether to use two backbones and fuse their features91 image_sizes: List of image sizes for each backbone92 timm_model_ids: List of TIMM model IDs to use for each backbone93 timm_override_act_layers: List of activation layer overrides for each backbone94 """95 super().__init__()96 self.use_fused_vision_backbone = use_fused_vision_backbone97 self.num_images_in_input = 1 # Default value, can be overridden later98 99 # Validate number of (fused) vision backbones100 if len(timm_model_ids) > 2:101 raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")102 103 # Create primary featurizer104 self.featurizer = self._create_featurizer(105 model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]106 )107 self.embed_dim = self.featurizer.embed_dim108 109 # Create secondary featurizer if using fused backbone110 if self.use_fused_vision_backbone:111 self.fused_featurizer = self._create_featurizer(112 model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]113 )114 self.embed_dim += self.fused_featurizer.embed_dim115 116 # Patch LayerScale modules for HF compatibility117 self._patch_layer_scales()118 119 120 def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:121 """122 Create a TIMM-based featurizer model with appropriate configurations.123 124 Args:125 model_id: The TIMM model ID to load126 img_size: Input image size for the model127 act_layer: Override for the activation layer type128 129 Returns:130 A configured featurizer model131 """132 featurizer = timm.create_model(133 model_id,134 pretrained=False,135 num_classes=0,136 img_size=img_size,137 act_layer=act_layer,138 )139 140 # Monkey-patch the forward function to extract the second-to-last layer features141 num_blocks = len(featurizer.blocks)142 featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))143 144 return featurizer145 146 147 def _patch_layer_scales(self) -> None:148 """149 Patch all LayerScale modules to be compatible with HF's parameter naming.150 151 HF Transformers overwrites parameters with names containing 'gamma',152 so we need to rename and modify the forward method.153 """154 # Patch primary featurizer155 for module in self.featurizer.modules():156 if isinstance(module, LayerScale):157 ls_apply_patch(module)158 159 # Patch secondary featurizer if it exists160 if self.use_fused_vision_backbone:161 for module in self.fused_featurizer.modules():162 if isinstance(module, LayerScale):163 ls_apply_patch(module)164 165 166 def get_num_patches(self) -> int:167 """168 Returns the number of vision patches output by the vision backbone.169 170 Returns:171 Number of patches per image172 """173 return self.featurizer.patch_embed.num_patches174 175 176 def get_num_images_in_input(self) -> int:177 """178 Returns the number of input images for the vision backbone.179 180 Returns:181 Number of images expected in the input182 """183 return self.num_images_in_input184 185 186 def set_num_images_in_input(self, num_images_in_input: int) -> None:187 """188 Sets the number of input images for the vision backbone.189 190 Args:191 num_images_in_input: Number of images to expect in the input192 """193 self.num_images_in_input = num_images_in_input194 195 196 def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:197 """198 Implements the forward pass for the vision backbone.199 200 If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features201 (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).202 203 Args:204 pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).205 """206 if self.num_images_in_input == 1:207 if not self.use_fused_vision_backbone:208 return self.featurizer(pixel_values)209 210 # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack211 img, img_fused = torch.split(pixel_values, [3, 3], dim=1)212 patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)213 214 return torch.cat([patches, patches_fused], dim=2)215 216 else:217 assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"218 219 # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)220 images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)221 222 # Process each image and collect patches223 all_patches = []224 for img in images:225 # Split each image further into two stacks of channels (each with 3 channels)226 img_regular, img_fused = torch.split(img, [3, 3], dim=1)227 228 # Get patches from both SigLIP and DINOv2 vision transformers229 patches = self.featurizer(img_regular)230 patches_fused = self.fused_featurizer(img_fused)231 232 # Concatenate SigLIP and DINOv2 patches along the hidden dimension233 combined_patches = torch.cat([patches, patches_fused], dim=2)234 all_patches.append(combined_patches)235 236 # Concatenate all patches along the patch dimension237 return torch.cat(all_patches, dim=1)238 239 240 241# === Prismatic Projector (nn.Module) Definitions ===242class PrismaticProjector(nn.Module):243 def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:244 super().__init__()245 self.use_fused_vision_backbone = use_fused_vision_backbone246 self.vision_dim, self.llm_dim = vision_dim, llm_dim247 248 # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!249 if not self.use_fused_vision_backbone:250 self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)251 self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)252 self.act_fn1 = nn.GELU()253 else:254 initial_projection_dim = 4 * vision_dim255 self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)256 self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)257 self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)258 self.act_fn1 = nn.GELU()259 self.act_fn2 = nn.GELU()260 261 def forward(self, img_patches: torch.Tensor) -> torch.Tensor:262 if not self.use_fused_vision_backbone:263 projected_features = self.fc1(img_patches)264 projected_features = self.act_fn1(projected_features)265 projected_features = self.fc2(projected_features)266 else:267 projected_features = self.fc1(img_patches)268 projected_features = self.act_fn1(projected_features)269 projected_features = self.fc2(projected_features)270 projected_features = self.act_fn2(projected_features)271 projected_features = self.fc3(projected_features)272 273 return projected_features274 275 276 277# === Main HF Class Definitions ===278@dataclass279class PrismaticCausalLMOutputWithPast(ModelOutput):280 """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""281 282 loss: Optional[torch.FloatTensor] = None283 logits: torch.FloatTensor = None284 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None285 hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None286 attentions: Optional[Tuple[torch.FloatTensor]] = None287 288 # Additions for VLMs289 projector_features: Optional[torch.FloatTensor] = None290 291 292 293class PrismaticPreTrainedModel(PreTrainedModel):294 config_class: PretrainedConfig = PrismaticConfig295 base_model_prefix: str = "model"296 supports_gradient_checkpointing: bool = True297 298 _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]299 _skip_keys_device_placement: str = "past_key_values"300 _supports_flash_attn_2: bool = True301 302 def _init_weights(self, module: nn.Module) -> None:303 # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!304 # => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at305 # https://github.com/TRI-ML/prismatic-vlms306 std = (307 self.config.initializer_range308 if hasattr(self.config, "initializer_range")309 else self.config.text_config.initializer_range310 )311 312 if hasattr(module, "class_embedding"):313 module.class_embedding.data.normal_(mean=0.0, std=std)314 315 if isinstance(module, (nn.Linear, nn.Conv2d)):316 module.weight.data.normal_(mean=0.0, std=std)317 if module.bias is not None:318 module.bias.data.zero_()319 elif isinstance(module, nn.Embedding):320 module.weight.data.normal_(mean=0.0, std=std)321 if module.padding_idx is not None:322 module.weight.data[module.padding_idx].zero_()323 324 @property325 def _supports_sdpa(self) -> bool:326 """Check LLM supports SDPA Attention"""327 return self.language_model._supports_sdpa328 329 330 331class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):332 def __init__(self, config: PrismaticConfig) -> None:333 super().__init__(config)334 335 # [Validation] Lightweight Validate on `config` Fields + Dependency Versions336 if config.use_fused_vision_backbone is None:337 raise ValueError("Missing config field `use_fused_vision_backbone`")338 339 if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:340 raise NotImplementedError(341 "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "342 "if you urgently need support for latest TIMM versions."343 )344 345 if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):346 logger.warning(347 f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "348 f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "349 f"there might be inference-time regressions due to dependency changes. If in doubt, please"350 f"use the above versions."351 )352 353 # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)354 self.vision_backbone = PrismaticVisionBackbone(355 config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers356 )357 358 # Create Multimodal Projector359 self.projector = PrismaticProjector(360 config.use_fused_vision_backbone,361 vision_dim=self.vision_backbone.embed_dim,362 llm_dim=config.text_config.hidden_size,363 )364 365 # Instantiate LLM Backbone366 self.language_model = AutoModelForCausalLM.from_config(367 config.text_config, attn_implementation=config._attn_implementation368 )369 370 self.vocab_size = config.text_config.vocab_size371 self.pad_token_id = config.pad_token_id372 self.llm_dim = config.text_config.hidden_size373 374 #Action query token375 self.action_queries = nn.Embedding(NUM_TOKENS, self.llm_dim)376 self.action_queries.weight.data.zero_()377 378 # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing379 self.post_init()380 381 # === `PreTrainedModel` Boilerplate ===382 def get_input_embeddings(self) -> nn.Module:383 return self.language_model.get_input_embeddings()384 def set_version(self, version: str):385 self.version = version386 return self.version387 388 389 def set_input_embeddings(self, value: nn.Module) -> None:390 self.language_model.set_input_embeddings(value)391 392 def get_output_embeddings(self) -> nn.Module:393 return self.language_model.get_output_embeddings()394 395 def set_output_embeddings(self, new_embeddings: nn.Module) -> None:396 self.language_model.set_output_embeddings(new_embeddings)397 398 def get_decoder(self) -> nn.Module:399 return self.language_model.get_decoder()400 401 def set_decoder(self, decoder: nn.Module) -> None:402 self.language_model.set_decoder(decoder)403 404 def tie_weights(self) -> None:405 self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)406 407 def resize_token_embeddings(408 self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None409 ) -> nn.Embedding:410 updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)411 412 # Update config/instance variables413 self.config.text_config.vocab_size = updated_embeddings.num_embeddings414 self.vocab_size = updated_embeddings.num_embeddings415 416 return updated_embeddings417 418 def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):419 """420 Replace embeddings in input_embeddings at positions where all_actions_mask is True421 with embeddings from noisy_action_features, using vectorized operations.422 423 Args:424 input_embeddings: Tensor of shape (B, S, D)425 all_actions_mask: Boolean tensor of shape (B, S)426 noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample427 428 Returns:429 Modified input_embeddings tensor430 """431 # Clone input to avoid modifying the original tensor432 new_input_embeddings = input_embeddings.clone()433 434 # Create a tensor with the same shape of input_embeddings to hold the noisy action features435 repositioned_noisy_action_features = torch.zeros_like(input_embeddings)436 437 # Create batch indices for splicing438 batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)439 batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])440 441 # Get indices where mask is True for each sample442 masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])443 444 # Move the noisy action features into their correct positions445 # print(noisy_action_features.size())446 447 repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features448 449 # Combine original input embeddings and noisy action embeddings using the mask450 new_input_embeddings = torch.where(451 all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings452 )453 454 return new_input_embeddings455 456 def _process_action_masks(self, labels):457 """Helper to get action masks from labels"""458 current_action_mask = get_current_action_mask(labels)459 next_actions_mask = get_next_actions_mask(labels)460 all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)461 return all_actions_mask462 463 def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):464 """Process vision features with optional FiLM conditioning"""465 if use_film:466 # FiLM: Infuse language inputs into visual features467 patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D)468 else:469 patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)470 471 # Project patch embeddings into language embedding space472 return self.projector(patch_features)473 474 def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):475 """Process proprioceptive features and append to vision features"""476 if proprio_projector is not None and proprio is not None:477 # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)478 # proprio: (bsz, proprio_dim) or (propro_dim,)479 proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)480 proprio_features = proprio_projector(proprio) # (bsz, llm_dim)481 proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)482 # For simplicity, just append proprio token to the end of projected vision patch tokens483 return torch.cat((projected_patch_embeddings, proprio_features), dim=1)484 return projected_patch_embeddings485 486 def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):487 """Build multimodal embeddings and attention mask"""488 # Update attention mask489 490 projected_patch_attention_mask = None491 if attention_mask is not None:492 projected_patch_attention_mask = torch.full(493 (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),494 fill_value=True,495 dtype=attention_mask.dtype,496 device=attention_mask.device,497 )498 499 # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)500 multimodal_embeddings = torch.cat(501 [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1502 )503 504 multimodal_attention_mask = None505 if attention_mask is not None:506 multimodal_attention_mask = torch.cat(507 [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1508 )509 510 return multimodal_embeddings, multimodal_attention_mask511 512 def _build_multimodal_labels(self, labels, projected_patch_embeddings):513 """Build multimodal labels with IGNORE_INDEX for patch embeddings"""514 if labels is not None:515 projected_patch_labels = torch.full(516 (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),517 fill_value=IGNORE_INDEX,518 dtype=labels.dtype,519 device=labels.device,520 )521 return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)522 return None523 524 # === Core Prismatic VLM `forward()` Logic ===525 def forward(526 self,527 input_ids: Optional[torch.LongTensor] = None,528 attention_mask: Optional[torch.Tensor] = None,529 pixel_values: Optional[torch.FloatTensor] = None,530 labels: Optional[torch.LongTensor] = None,531 inputs_embeds: Optional[torch.FloatTensor] = None,532 past_key_values: Optional[List[torch.FloatTensor]] = None,533 use_cache: Optional[bool] = None,534 output_attentions: Optional[bool] = None,535 output_hidden_states: Optional[bool] = None,536 output_projector_features: Optional[bool] = None,537 return_dict: Optional[bool] = None,538 proprio=None,539 proprio_projector=None,540 noisy_actions=None,541 noisy_action_projector=None,542 diffusion_timestep_embeddings=None,543 use_film: bool = False,544 ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:545 """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""546 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions547 output_hidden_states = (548 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states549 )550 output_projector_features = output_projector_features if output_projector_features is not None else False551 return_dict = return_dict if return_dict is not None else self.config.use_return_dict552 553 # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)554 use_cache = use_cache and not self.training555 556 # Instantiate Placeholder for Projector Features557 projected_patch_embeddings = None558 559 # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===560 if input_ids.shape[1] == 1:561 assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"562 assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"563 assert labels is None, "Unexpected key `labels` provided during cached generation!"564 565 language_model_output = self.language_model(566 input_ids=input_ids,567 attention_mask=None,568 position_ids=None,569 past_key_values=past_key_values,570 inputs_embeds=None,571 labels=None,572 use_cache=use_cache,573 output_attentions=output_attentions,574 output_hidden_states=output_hidden_states,575 return_dict=return_dict,576 )577 578 # === Handle Unimodal Forward ===579 elif pixel_values is None:580 assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"581 assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"582 583 language_model_output = self.language_model(584 input_ids=input_ids,585 attention_mask=attention_mask,586 position_ids=None,587 past_key_values=None,588 inputs_embeds=None,589 labels=labels,590 use_cache=use_cache,591 output_attentions=output_attentions,592 output_hidden_states=output_hidden_states,593 return_dict=return_dict,594 )595 596 # === Handle Multimodal Forward ===597 elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):598 assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"599 600 # Get input embeddings (from language model embeddings)601 input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)602 603 604 # Extract action masks605 all_actions_mask = self._process_action_masks(labels)606 607 # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)608 609 # print(input_embeddings[~all_actions_mask].size())610 language_embeddings = input_embeddings[~all_actions_mask].reshape(611 input_embeddings.shape[0], -1, input_embeddings.shape[2]612 ) # (B, lang_seq_len, llm_dim)613 614 # Get visual features615 projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)616 617 # Process action embeddings618 if noisy_actions is not None:619 620 621 action_queries = self.action_queries.weight # (1, h)622 action_queries = action_queries.view(1, action_queries.shape[0], action_queries.shape[1]).repeat(input_embeddings.shape[0], 1, 1) # (b, chunk_size, h)623 all_actions_mask = self._process_action_masks(labels)624 input_embeddings = self._replace_input_embeddings(625 input_embeddings, all_actions_mask, action_queries)626 627 628 else:629 action_queries = self.action_queries.weight # (1, h)630 action_queries = action_queries.view(1, action_queries.shape[0], action_queries.shape[1]).repeat(input_embeddings.shape[0], 1, 1) # (b, chunk_size, h)631 all_actions_mask = self._process_action_masks(labels)632 input_embeddings = self._replace_input_embeddings(633 input_embeddings, all_actions_mask, action_queries)634 635 # Build multimodal embeddings & attention mask636 multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(637 input_embeddings, projected_patch_embeddings, attention_mask638 )639 640 # Build labels for multimodal sequence if needed641 multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)642 643 # Dispatch to language model644 language_model_output = self.language_model(645 input_ids=None,646 attention_mask=multimodal_attention_mask,647 position_ids=None,648 past_key_values=None,649 inputs_embeds=multimodal_embeddings,650 labels=None,651 use_cache=use_cache,652 output_attentions=output_attentions,653 output_hidden_states=output_hidden_states,654 return_dict=return_dict,655 ) 656 657 # === Otherwise =>> Assume Invalid! ===658 elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):659 raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")660 661 else:662 raise ValueError(663 "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"664 f"=> `input_ids` = {input_ids is not None}\n"665 f"=> `attention_mask` = {attention_mask is not None}\n"666 f"=> `pixel_values` = {pixel_values is not None}\n"667 f"=> `labels` = {labels is not None}\n"668 f"=> `input_embeds` = {inputs_embeds is not None}\n"669 f"=> `past_key_values` = {past_key_values is not None}\n"670 f"=> `use_cache` = {use_cache}"671 )672 673 # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)674 if not return_dict:675 if output_projector_features and (projected_patch_embeddings is not None):676 return *language_model_output, projected_patch_embeddings677 678 return language_model_output679 680 return PrismaticCausalLMOutputWithPast(681 loss=language_model_output.loss,682 past_key_values=language_model_output.past_key_values,683 hidden_states=language_model_output.hidden_states,684 attentions=language_model_output.attentions,685 projector_features=projected_patch_embeddings,686 )687 688 689 # === GenerationMixin Methods ===690 def prepare_inputs_for_generation(691 self,692 input_ids: Optional[torch.Tensor] = None,693 past_key_values: Optional[List[torch.FloatTensor]] = None,694 inputs_embeds: Optional[torch.FloatTensor] = None,695 pixel_values: Optional[torch.FloatTensor] = None,696 attention_mask: Optional[torch.Tensor] = None,697 **kwargs: str,698 ) -> Dict[str, torch.Tensor]:699 """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""700 if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (701 (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)702 ):703 raise ValueError("Generation with batch size > 1 is not currently supported!")704 705 # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens706 if past_key_values is not None:707 input_ids = input_ids[:, -1:]708 709 # If `input_embeds` are passed, we only want to use them in the 1st generation step710 if inputs_embeds is not None and past_key_values is None:711 model_inputs = {"input_embeds": inputs_embeds}712 else:713 model_inputs = {"input_ids": input_ids}714 715 # Make sure `pixel_values` are preserved in `model_inputs`716 model_inputs.update(717 {718 "attention_mask": attention_mask,719 "pixel_values": pixel_values,720 "past_key_values": past_key_values,721 "use_cache": kwargs.get("use_cache"),722 }723 )724 725 return model_inputs726 727 # Defer to Language Model (all handle this differently, with different return types)728 def _reorder_cache(self, *args, **kwargs) -> Any:729 return self.language_model._reorder_cache(*args, **kwargs)730 731 732 733class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):734 config_class: PretrainedConfig = OpenVLAConfig735 736 def __init__(self, config: OpenVLAConfig) -> None:737 super().__init__(config)738 self.norm_stats = config.norm_stats739 740 741 # Compute action bins742 self.bins = np.linspace(-1, 1, config.n_action_bins)743 self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0744 745 # Compute vocab size for de-tokenization -- revert added "multiple of"746 self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of747 748 def _prepare_input_for_action_prediction(self, input_ids, attention_mask):749 """Prepares input for action prediction by adding necessary tokens"""750 # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens751 placeholder_action_token_ids = (752 torch.ones((input_ids.shape[0], NUM_TOKENS)).to(input_ids.device).to(input_ids.dtype)753 )754 input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)755 756 # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)757 stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX758 input_ids = torch.cat([input_ids, stop_token_id], dim=-1)759 760 # Extend the attention mask to fit the new shape of input761 # Note: Only batch size == 1 supported right now762 mask_extension = (763 torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))764 .to(attention_mask.device)765 .to(attention_mask.dtype)766 )767 attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)768 769 return input_ids, attention_mask770 771 def _prepare_labels_for_action_prediction(self, labels, input_ids):772 """Creates labels tensor for action prediction if not provided"""773 # Extend labels tensor with fake action labels774 ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1775 labels_extension = (776 torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)777 * ARBITRARY_ACTION_TOKEN_IDX778 )779 labels = torch.cat([labels, labels_extension], dim=-1)780 781 # Replace last label token with stop token782 labels[:, -1] = STOP_INDEX783 784 return labels785 786 def _unnormalize_actions(self, normalized_actions, unnorm_key=None):787 """Unnormalize actions using dataset statistics"""788 action_norm_stats = self.get_action_stats(unnorm_key)789 790 if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:791 mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))792 action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])793 elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:794 mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))795 action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])796 else:797 raise ValueError("Unsupported action/proprio normalization type detected!")798 799 actions = np.where(800 mask,801 0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,802 normalized_actions,803 )804 805 return actions806 807 808 def _regression_or_discrete_prediction(809 self,810 input_embeddings,811 all_actions_mask,812 projected_patch_embeddings,813 attention_mask,814 labels,815 NUM_PATCHES,816 NUM_PROMPT_TOKENS,817 action_head=None,818 proprio=None,819 proprio_projector=None,820 ):821 """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""822 823 action_queries = self.action_queries.weight # (1, h)824 action_queries = action_queries.view(1, action_queries.shape[0], action_queries.shape[1]).repeat(input_embeddings.shape[0], 1, 1) # (b, chunk_size, h)825 # Replace action token embeddings with noisy action embeddings826 input_embeddings = self._replace_input_embeddings(input_embeddings.clone(), all_actions_mask, action_queries)827 828 # Build multimodal embeddings and attention mask829 multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(830 input_embeddings, projected_patch_embeddings, attention_mask831 )832 833 # Forward pass through language model834 language_model_output = self.language_model(835 input_ids=None,836 attention_mask=multimodal_attention_mask,837 position_ids=None,838 past_key_values=None,839 inputs_embeds=multimodal_embeddings,840 labels=None,841 use_cache=None,842 output_attentions=False,843 output_hidden_states=True,844 return_dict=True,845 )846 847 # Extract hidden states for action tokens848 multi_layer_hidden_states = []849 850 for item in language_model_output.hidden_states[0:]:851 # last_hidden_states = output.hidden_states[-1] # (B, seq_len, D)852 # Get hidden states for text portion of prompt+response (after the vision patches)853 text_hidden_states = item854 # Get hidden states for action portion of response855 actions_hidden_states = text_hidden_states[:, NUM_PATCHES+ NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_TOKENS, :,].reshape(1, 1, NUM_TOKENS, -1).to(torch.bfloat16)856 857 batch_size = item.shape[0]858 task_latten_states = item[:, :NUM_PATCHES].reshape(batch_size, 1, NUM_PATCHES , -1)859 all_hidden_states = torch.cat((task_latten_states, actions_hidden_states),2)860 multi_layer_hidden_states.append(all_hidden_states)861 862 multi_layer_hidden_states = torch.cat(multi_layer_hidden_states, dim = 1)863 864 865 # Handle different prediction methods866 if action_head is not None:867 # L1 regression prediction868 normalized_actions = action_head.predict_action(multi_layer_hidden_states,869 proprio=proprio,870 proprio_projector=proprio_projector)871 normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)872 normalized_actions = normalized_actions.float().cpu().detach().numpy()873 else:874 # Discrete token-based prediction875 predicted_action_token_ids = (876 language_model_output.logits[877 :,878 NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,879 ]880 .argmax(dim=2)881 .cpu()882 .numpy()883 )884 discretized_actions = self.vocab_size - predicted_action_token_ids885 discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)886 normalized_actions = self.bin_centers[discretized_actions]887 normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)888 889 return normalized_actions, actions_hidden_states890 891 892 def predict_action(893 self,894 input_ids: Optional[torch.LongTensor] = None,895 unnorm_key: Optional[str] = None,896 proprio=None,897 proprio_projector=None,898 action_head=None,899 noisy_action_projector=None,900 use_film: bool = False,901 **kwargs: str,902 ) -> np.ndarray:903 """Predict actions from input sequence, with options for different prediction methods.904 905 Args:906 input_ids: Input token ids907 unnorm_key: Key for unnormalization statistics908 proprio: Proprioceptive features909 proprio_projector: Projector for proprioceptive features910 action_head: Optional head for L1 regression or diffusion-based prediction911 noisy_action_projector: Projector for noisy actions in diffusion-based prediction912 use_film: Whether to use FiLM conditioning913 **kwargs: Additional arguments including pixel_values and attention_mask914 915 Returns:916 Tuple of (unnormalized_actions, action_hidden_states)917 """918 919 pixel_values = kwargs["pixel_values"] # [1, 12, 224, 224]920 attention_mask = kwargs["attention_mask"] # 921 922 # Create fake labels tensor (needed for action mask)923 labels = input_ids.clone()924 labels[:] = IGNORE_INDEX925 926 # Get number of tokens in prompt (excluding the start token)927 NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token928 929 # Prepare inputs by adding necessary tokens930 input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)931 932 # Update labels tensor for action mask computation later933 labels = self._prepare_labels_for_action_prediction(labels, input_ids)934 935 # Get input embeddings and action masks936 input_embeddings = self.get_input_embeddings()(input_ids)937 all_actions_mask = self._process_action_masks(labels)938 939 # Extract language embeddings940 language_embeddings = input_embeddings[~all_actions_mask].reshape(941 input_embeddings.shape[0], -1, input_embeddings.shape[2]942 )943 944 # Process vision features945 projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)946 947 # Add proprioceptive features if provided948 use_proprio = proprio_projector is not None and proprio is not None949 if use_proprio:950 proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)951 952 # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)953 NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()954 955 # Run regression or discrete token-based prediction956 normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(957 input_embeddings,958 all_actions_mask,959 projected_patch_embeddings,960 attention_mask,961 labels,962 NUM_PATCHES,963 NUM_PROMPT_TOKENS,964 action_head=action_head,965 proprio=proprio, # [8]966 proprio_projector=proprio_projector,967 )968 969 # Unnormalize predicted actions970 actions = self._unnormalize_actions(normalized_actions, unnorm_key)971 972 return actions, actions_hidden_states973 974 975 976 @staticmethod977 def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:978 """Validate and resolve the unnormalization key for action statistics"""979 if unnorm_key is None:980 assert len(norm_stats) == 1, (981 f"Your model was trained on more than one dataset, "982 f"please pass a `unnorm_key` from the following options to choose the statistics "983 f"used for un-normalizing actions: {norm_stats.keys()}"984 )985 unnorm_key = next(iter(norm_stats.keys()))986 987 assert unnorm_key in norm_stats, (988 f"The `unnorm_key` you chose is not in the set of available dataset statistics, "989 f"please choose from: {norm_stats.keys()}"990 )991 return unnorm_key992 993 def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:994 """Get the dimensionality of the policy's action space."""995 unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)996 return len(self.norm_stats[unnorm_key]["action"]["min"])997 998 def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:999 """Get all the logged statistics for the given dataset."""1000 unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)1001 return self.norm_stats[unnorm_key]["action"]1002