bitvla_for_action_prediction.py

from transformers import LlavaForConditionalGeneration,PretrainedConfig
from configuration_bit_vla import Bitvla_Config
import numpy as np
import torch
from prismatic.vla.constants import (
    ACTION_DIM,
    ACTION_PROPRIO_NORMALIZATION_TYPE,
    NUM_ACTIONS_CHUNK,
    NormalizationType,
)
from typing import Optional, Dict, Any,List,Tuple

from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast

from prismatic.training.train_utils import (
    get_current_action_mask,
    get_next_actions_mask,
)


class BitVLAForActionPrediction(LlavaForConditionalGeneration):
    config_class: PretrainedConfig = Bitvla_Config

    def __init__(self, config) -> None:
        super().__init__(config)
        self.norm_stats = config.norm_stats

        # Compute action bins
        self.bins = np.linspace(-1, 1, config.n_action_bins)
        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0

        self.vocab_size = self.config.vocab_size

    def set_constant(self, image_token_idx, proprio_pad_idx, ignore_idx, action_token_begin_idx, stop_index):
        self.image_token_idx = image_token_idx
        self.proprio_pad_idx = proprio_pad_idx
        self.action_token_begin_idx = action_token_begin_idx
        self.stop_index = stop_index
        self.ignore_idx = ignore_idx
    
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_projector_features: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        proprio=None,
        proprio_projector=None,
        cache_position: Optional[torch.LongTensor] = None,
        vision_feature_layer=None,
        vision_feature_select_strategy=None,
        ) -> Tuple[int, LlavaCausalLMOutputWithPast]:
        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        output_projector_features = output_projector_features if output_projector_features is not None else False
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
        use_cache = use_cache and not self.training

        batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] # type: ignore

        # === Handle Multimodal Forward ===
        if (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"

            # Get input embeddings
            inputs_embeds = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
            
            # change the vision padding to the real vision tokens
            if pixel_values is not None:
                vision_feature_layer = (
                    vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
                )
                vision_feature_select_strategy = (
                    vision_feature_select_strategy
                    if vision_feature_select_strategy is not None
                    else self.config.vision_feature_select_strategy
                )
                # pixel_values: b,num_images,c,h,w
                # for each image, we do self.get_image_features
                # then we concat the features of all images
                # pixel_values: (b,num_images,c,h,w) --> (b*num_images,c,h,w)
                b, num_images, c, h, w = pixel_values.shape
                pixel_values = pixel_values.view(-1, c, h, w)  # (b*num_images,c,h,w)
                image_embeds = self.get_image_features(
                    pixel_values = pixel_values,
                    vision_feature_layer = vision_feature_layer,
                    vision_feature_select_strategy = vision_feature_select_strategy,
                )
                
                # image_features: (b*num_images,seq_len,patch_size) --> (b*num_images*seq_len,patch_size)
                image_embeds = image_embeds.view(-1,image_embeds.shape[-1])
                n_image_tokens = (input_ids == self.image_token_idx).sum().item()
                n_image_features = image_embeds.shape[0]
                if n_image_tokens != n_image_features:
                    raise ValueError(
                        f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
                    )

                mask = input_ids == self.image_token_idx
                mask_unsqueezed = mask.unsqueeze(-1)
                mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
                image_mask = mask_expanded.to(inputs_embeds.device)

                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
                
                
            # change the proprio padding to the real proprio tokens
            if proprio_projector is not None and proprio is not None:
                # proprio: (bsz, proprio_dim) or (propro_dim,)
                proprio = proprio.reshape(batch_size, -1)  # (bsz, proprio_dim)
                proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
                proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
                #（bsz, 1, llm_dim) --> (bsz*1, llm_dim)
                proprio_features = proprio_features.view(-1, proprio_features.shape[-1])
                n_proprio_tokens = (input_ids == self.proprio_pad_idx).sum().item()
                n_proprio_features = proprio_features.shape[0]
                if n_proprio_tokens != n_proprio_features:
                    raise ValueError(
                        f"Proprio features and proprio tokens do not match: tokens: {n_proprio_tokens}, features {n_proprio_features}"
                    )
                
                mask = input_ids == self.proprio_pad_idx
                mask_unsqueezed = mask.unsqueeze(-1)
                mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
                proprio_mask = mask_expanded.to(inputs_embeds.device)
                
                proprio_features = proprio_features.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(proprio_mask, proprio_features)
                
            
            # Extract action masks
            # Action tokens are those in labels that are not ignore, not newline, and not end-of-sequence tokens
            all_actions_mask = (labels != self.ignore_idx) & (labels != self.stop_index)
          
            # Replace the embeddings of the action tokens with zeros
            # (Later on, the positional embeddings will be added to them)
            all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
            inputs_embeds = inputs_embeds * ~all_actions_mask 
            outputs = LlavaForConditionalGeneration.forward(
                self,
                input_ids = None,
                attention_mask=attention_mask,
                position_ids=None,
                pixel_values=None,
                labels=labels,
                inputs_embeds=inputs_embeds,
                past_key_values=None,
                use_cache=None,
                output_attentions=False,
                output_hidden_states=True,
                return_dict=True,
            )
        # === Otherwise =>> Assume Invalid! ===
        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")

        else:
            raise ValueError(
                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
                f"=> `input_ids` = {input_ids is not None}\n"
                f"=> `attention_mask` = {attention_mask is not None}\n"
                f"=> `pixel_values` = {pixel_values is not None}\n"
                f"=> `labels` = {labels is not None}\n"
                f"=> `input_embeds` = {inputs_embeds is not None}\n"
                f"=> `past_key_values` = {past_key_values is not None}\n"
                f"=> `use_cache` = {use_cache}"
            )

        return outputs

    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
        """Prepares input for action prediction by adding necessary tokens"""
        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
        placeholder_action_token_ids = (
            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
        )
        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)

        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * self.stop_index
        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)

        # Extend the attention mask to fit the new shape of input
        # Note: Only batch size == 1 supported right now
        mask_extension = (
            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
            .to(attention_mask.device)
            .to(attention_mask.dtype)
        )
        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)

        return input_ids, attention_mask

    def _prepare_labels_for_action_prediction(self, labels, input_ids):
        """Creates labels tensor for action prediction if not provided"""
        # Extend labels tensor with fake action labels
        ARBITRARY_ACTION_TOKEN_IDX = self.action_token_begin_idx + 1
        labels_extension = (
            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
            * ARBITRARY_ACTION_TOKEN_IDX
        )
        labels = torch.cat([labels, labels_extension], dim=-1)

        # Replace last label token with stop token
        labels[:, -1] = self.stop_index

        return labels

    def _process_action_masks(self, labels):
        """Helper to get action masks from labels"""
        current_action_mask = get_current_action_mask(labels,ignore_index=self.ignore_idx,action_token_begin_idx=self.action_token_begin_idx)
        next_actions_mask = get_next_actions_mask(labels,ignore_index=self.ignore_idx,action_token_begin_idx=self.action_token_begin_idx)
        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
        return all_actions_mask
        
    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
        """Unnormalize actions using dataset statistics"""
        action_norm_stats = self.get_action_stats(unnorm_key)

        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
        else:
            raise ValueError("Unsupported action/proprio normalization type detected!")

        actions = np.where(
            mask,
            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
            normalized_actions,
        )

        return actions
    
    def _regression_or_discrete_prediction(
        self,
        input_ids,
        input_embeddings,
        all_actions_mask,
        attention_mask,
        labels,
        action_head=None,
        pixel_values = None,
    ):
        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
        # Zero out action token embeddings
        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
        input_embeddings = input_embeddings * ~all_actions_mask

        llava_output = LlavaForConditionalGeneration.forward(
                self,
                input_ids = None,
                attention_mask=attention_mask,
                position_ids=None,
                pixel_values=None,
                labels=None,
                inputs_embeds=input_embeddings,
                past_key_values=None,
                use_cache=None,
                output_attentions=False,
                output_hidden_states=True,
                return_dict=True,
            )
        all_actions_mask = self._process_action_masks(labels[:,1:])
        # Extract hidden states for action tokens
        last_hidden_states = llava_output.hidden_states[-1]  # (B, seq_len, D)
        last_hidden_states = last_hidden_states[:, : -1, :]  # (B, act_chunk_len, D)
        # Use the action mask to extract the hidden states of the actions
        actions_hidden_states = last_hidden_states[all_actions_mask.squeeze(-1)].unsqueeze(0) # (B, act_chunk_len, D)

        # Handle different prediction methods
        if action_head is not None:
            # L1 regression prediction
            normalized_actions = action_head.predict_action(actions_hidden_states)
            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
            normalized_actions = normalized_actions.float().cpu().detach().numpy()
        else:
            # Discrete token-based prediction
            predicted_action_token_ids = (
                llava_output.logits[all_actions_mask.squeeze(-1)].unsqueeze(0)
                .argmax(dim=2)
                .cpu()
                .numpy()
            )
            # FIXME: We do not support discrete action prediction right now
            # It seems that vocab_size here is not correct. This should be the dimension of the logit layer, which is actually larger than the vocab_size in the tokenizer. What we actually need here is the vocab_size from the tokenizer.
            discretized_actions = self.vocab_size - predicted_action_token_ids
            discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
            normalized_actions = self.bin_centers[discretized_actions]
            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)

        return normalized_actions, actions_hidden_states
    
    def predict_action(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        unnorm_key: Optional[str] = None,
        proprio=None,
        proprio_projector=None,
        action_head=None,
        vision_feature_layer=None,
        vision_feature_select_strategy=None,
        **kwargs: str,
    ) -> np.ndarray:
        """Predict actions from input sequence, with options for different prediction methods.

        Args:
            input_ids: Input token ids
            unnorm_key: Key for unnormalization statistics
            proprio: Proprioceptive features
            proprio_projector: Projector for proprioceptive features
            action_head: Optional head for L1 regression prediction
            **kwargs: Additional arguments including pixel_values and attention_mask

        Returns:
            Tuple of (unnormalized_actions, action_hidden_states)
        """
        pixel_values = kwargs["pixel_values"]
        attention_mask = kwargs["attention_mask"]

        # Create fake labels tensor (needed for action mask)
        labels = input_ids.clone()
        labels[:] = self.ignore_idx

        # Prepare inputs by adding necessary tokens
        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)

        # Update labels tensor for action mask computation later
        labels = self._prepare_labels_for_action_prediction(labels, input_ids)

        # Get input embeddings and action masks
        input_embeddings = self.get_input_embeddings()(input_ids)
        all_actions_mask = self._process_action_masks(labels)
        
        # vision tokens
        if pixel_values is not None:
            vision_feature_layer = (
                vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
            )
            vision_feature_select_strategy = (
                vision_feature_select_strategy
                if vision_feature_select_strategy is not None
                else self.config.vision_feature_select_strategy
            )
            # pixel_values: b,num_images,c,h,w
            # for each image, we do self.get_image_features
            # then we concat the features of all images
            # pixel_values: (b,num_images,c,h,w) --> (b*num_images,c,h,w)
            b, num_images, c, h, w = pixel_values.shape
            pixel_values = pixel_values.view(-1, c, h, w)  # (b*num_images,c,h,w)
            image_embeds = self.get_image_features(
                pixel_values = pixel_values,
                vision_feature_layer = vision_feature_layer,
                vision_feature_select_strategy = vision_feature_select_strategy,
            )
            
            # image_features: (b*num_images,seq_len,patch_size) --> (b*num_images*seq_len,patch_size)
            image_embeds = image_embeds.view(-1,image_embeds.shape[-1])
            n_image_tokens = (input_ids == self.image_token_idx).sum().item()
            n_image_features = image_embeds.shape[0]
            if n_image_tokens != n_image_features:
                raise ValueError(
                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
                )

            mask = input_ids == self.image_token_idx
            mask_unsqueezed = mask.unsqueeze(-1)
            mask_expanded = mask_unsqueezed.expand_as(input_embeddings)
            image_mask = mask_expanded.to(input_embeddings.device)

            image_embeds = image_embeds.to(input_embeddings.device, input_embeddings.dtype)
            input_embeddings = input_embeddings.masked_scatter(image_mask, image_embeds)

        # Add proprioceptive features if provided
        use_proprio = proprio_projector is not None and proprio is not None
        if use_proprio:
            batch_size = input_ids.shape[0] if input_ids is not None else input_embeddings.shape[0] # type: ignore
            proprio = torch.Tensor(proprio).to(input_embeddings.device, dtype=input_embeddings.dtype)
            if proprio_projector is not None and proprio is not None:
                # proprio: (bsz, proprio_dim) or (propro_dim,)
                proprio = proprio.reshape(batch_size, -1)  # (bsz, proprio_dim)
                proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
                proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
                #（bsz, 1, llm_dim) --> (bsz*1, llm_dim)
                proprio_features = proprio_features.view(-1, proprio_features.shape[-1])
                n_proprio_tokens = (input_ids == self.proprio_pad_idx).sum().item()
                n_proprio_features = proprio_features.shape[0]
                if n_proprio_tokens != n_proprio_features:
                    raise ValueError(
                        f"Proprio features and proprio tokens do not match: tokens: {n_proprio_tokens}, features {n_proprio_features}"
                    )
                
                mask = input_ids == self.proprio_pad_idx
                mask_unsqueezed = mask.unsqueeze(-1)
                mask_expanded = mask_unsqueezed.expand_as(input_embeddings)
                proprio_mask = mask_expanded.to(input_embeddings.device)
                
                proprio_features = proprio_features.to(input_embeddings.device, input_embeddings.dtype)
                input_embeddings = input_embeddings.masked_scatter(proprio_mask, proprio_features)

        # Run regression or discrete token-based prediction
        normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
            input_ids,
            input_embeddings,
            all_actions_mask,
            attention_mask,
            labels,
            action_head,
            pixel_values,
        )

        # Unnormalize predicted actions
        actions = self._unnormalize_actions(normalized_actions, unnorm_key)

        return actions, actions_hidden_states

    @staticmethod
    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
        """Validate and resolve the unnormalization key for action statistics"""
        if unnorm_key is None:
            assert len(norm_stats) == 1, (
                f"Your model was trained on more than one dataset, "
                f"please pass a `unnorm_key` from the following options to choose the statistics "
                f"used for un-normalizing actions: {norm_stats.keys()}"
            )
            unnorm_key = next(iter(norm_stats.keys()))

        assert unnorm_key in norm_stats, (
            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
            f"please choose from: {norm_stats.keys()}"
        )
        return unnorm_key

    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
        """Get the dimensionality of the policy's action space."""
        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
        return len(self.norm_stats[unnorm_key]["action"]["min"])

    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
        """Get all the logged statistics for the given dataset."""
        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
        return self.norm_stats[unnorm_key]["action"]