import torch from torch import nn import copy def get_intermediate_size(hidden_dim, ffn_dim_multiplier=4, multiple_of=256): hidden_dim = int(2 * hidden_dim / 3) hidden_dim = int(ffn_dim_multiplier * hidden_dim) hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) return hidden_dim import torch.nn.functional as F # noqa: N812 import torch from typing import Optional,Callable,Dict,Any from torch import nn from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLAttention,apply_multimodal_rotary_pos_emb,eager_attention_forward,repeat_kv from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLTextConfig from transformers import Qwen2_5_VLTextModel,Qwen2_5_VLForConditionalGeneration from transformers.cache_utils import Cache from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.processing_utils import Unpack from transformers import AutoProcessor from einops import rearrange, repeat from qwen_vl_utils import process_vision_info import PIL import json import math import numpy as np from huggingface_hub import hf_hub_download def create_sinusoidal_pos_embedding( time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu" ): """Computes sine-cosine positional embedding vectors for scalar positions.""" if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") if time.ndim != 1: raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") dtype = torch.float32 fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) period = min_period * (max_period / min_period) ** fraction # Compute the outer product scaling_factor = 1.0 / period * 2 * math.pi sin_input = scaling_factor[None, :] * time[:, None] pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) return pos_emb def apply_rope(x, positions, max_wavelength=10_000): """ Applies RoPE positions [B, L] to x [B, L, H, D]. """ d_half = x.shape[-1] // 2 device = x.device dtype = x.dtype x = x.to(torch.float32) freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device) timescale = max_wavelength**freq_exponents radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32) radians = radians[..., None, :] sin = torch.sin(radians) # .to(dtype=dtype) cos = torch.cos(radians) # .to(dtype=dtype) x1, x2 = x.split(d_half, dim=-1) res = torch.empty_like(x) res[..., :d_half] = x1 * cos - x2 * sin res[..., d_half:] = x2 * cos + x1 * sin return res.to(dtype) def make_att_2d_masks(pad_masks, att_masks): """Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to setup several types of attention, for example: [[1 1 1 1 1 1]]: pure causal attention. [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour. [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a block can attend all previous blocks and all tokens on the same block. Args: input_mask: bool[B, N] true if its part of the input, false if padding. mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token. """ if att_masks.ndim != 2: raise ValueError(att_masks.ndim) if pad_masks.ndim != 2: raise ValueError(pad_masks.ndim) cumsum = torch.cumsum(att_masks, dim=1) att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] att_2d_masks = att_2d_masks & pad_2d_masks return att_2d_masks class Qwen2_5_VLMoTAttention(Qwen2_5_VLAttention): """ """ def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None): super().__init__(config,layer_idx) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC fill_kv_cache=True, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) #cos, sin = position_embeddings ## Since our action chunk is 1d time series, we do not need multimodal rope. Switch to normal rope instead #query_states, key_states = apply_multimodal_rotary_pos_emb( # query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] #) query_states = rearrange(query_states, 'b h s d -> b s h d') query_states = apply_rope(query_states,position_ids) query_states = rearrange(query_states, 'b s h d -> b h s d') key_states = rearrange(key_states, 'b h s d -> b s h d') key_states = apply_rope(key_states,position_ids) key_states = rearrange(key_states, 'b s h d -> b h s d') if use_cache: past_key_state = past_key_value[self.layer_idx][0] past_value_state = past_key_value[self.layer_idx][1] key_states = torch.cat([past_key_state, key_states], dim=2) # print(key_states.dtype) value_states = torch.cat( [past_value_state, value_states], dim=2 ) key_states = key_states.to(dtype=query_states.dtype) value_states = value_states.to(dtype=query_states.dtype) #print("New K shape",key_states.shape) #print("New V shape",value_states.shape) #if past_key_value is not None and not fill_kv_cache: ## Only update KV cache if fill_kv_cache is False #cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models # key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] #print("New query shape",query_states.shape) #attention_mask = torch.ones() ## I need to check if is_casual is default to True here. Is casual will automatically create an attention mask and I do not want that to happen. #print(position_ids) #print(attention_mask.shape) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, position_ids=position_ids, # pass positions for FA2 **kwargs, ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights from transformers.modeling_outputs import BaseModelOutputWithPast class Qwen2_5_VLAExpert(Qwen2_5_VLTextModel): def __init__(self,config): super().__init__(config) def forward(self, expert_attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, vlm_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs: Unpack[FlashAttentionKwargs],): 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 ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False if inputs_embeds is None: raise ValueError("You must specify exactly inputs_embeds") # torch.jit.trace() doesn't support cache objects in the output if vlm_key_values is None: raise ValueError("You must specify vlm_cache") hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers #position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) layer_outputs = decoder_layer( hidden_states, attention_mask=expert_attention_mask, position_ids=position_ids, past_key_value=vlm_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=None, **kwargs, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple( v for v in [hidden_states, vlm_key_values, all_hidden_states, all_self_attns] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=vlm_key_values, hidden_states=all_hidden_states, attentions=all_self_attns, ) import tensorflow as tf import dlimp as dl import PIL.Image as Image def resize_image(image1): #image1 = ds_combined[0]['observation.images.scene'] #image1 = image1.reshape(480,640,3) image1 = tf.cast(image1*255, dtype=tf.uint8) image1 = image1.numpy().transpose(1,2,0) image1 = dl.transforms.resize_image(image1, size=(224,224)) image1 = Image.fromarray(image1.numpy()) return image1 class VLAWithExpert(nn.Module): _ACTION_TOKEN_MIN = 151665 _ACTION_TOKEN_MAX = 153712 def __init__(self,config=None,device=None): super().__init__() self.vlm = Qwen2_5_VLForConditionalGeneration.from_pretrained( "declare-lab/nora-long", torch_dtype=torch.bfloat16, attn_implementation="sdpa", ) if config is not None: self.config = config else: self.config = {'max_action_dim':7,"max_state_dim":8} print("Loading expert model...") self.lm_expert_config = copy.deepcopy(self.vlm.config.text_config) #lm_expert_config = copy.deepcopy(model.config.text_config) self.processor = AutoProcessor.from_pretrained( "declare-lab/nora", trust_remote_code=True ) self.fast_tokenizer = fast_tokenizer = AutoProcessor.from_pretrained( "physical-intelligence/fast", trust_remote_code=True ) self.fast_tokenizer.action_dim = 7 self.fast_tokenizer.time_horizon = 5 hidden_size = self.lm_expert_config.hidden_size expert_width_multiplier = 0.375 self.lm_expert_config.hidden_size = int(hidden_size * expert_width_multiplier) # hidden_size // 2 self.lm_expert_config.intermediate_size = get_intermediate_size(int(hidden_size * expert_width_multiplier)) self.lm_expert_config.num_hidden_layers = self.vlm.config.num_hidden_layers self.lm_expert_config.num_attention_heads = 6 self.action_expert = Qwen2_5_VLAExpert._from_config(self.lm_expert_config,torch_dtype=torch.bfloat16) self.action_chunk_length = 5 self.device = self.vlm.device # Replace the action expert's attention layers self._replace_action_expert_attention() self.action_expert.embed_tokens = None self.vlm_kv_cache = None # self.state_proj = nn.Linear( # self.config['max_state_dim'], hidden_size # ) self.action_in_proj = nn.Linear(self.config['max_action_dim'],self.lm_expert_config.hidden_size) self.action_out_proj = nn.Linear(self.lm_expert_config.hidden_size, self.config['max_action_dim']) self.action_time_mlp_in = nn.Linear( self.lm_expert_config.hidden_size * 2, self.lm_expert_config.hidden_size ) self.action_time_mlp_out = nn.Linear( self.lm_expert_config.hidden_size, self.lm_expert_config.hidden_size ) self.state_emb = nn.Linear(self.config['max_action_dim'], self.lm_expert_config.hidden_size) self.device = self.vlm.device print(f"*** Loading normalization stats from HF Hub ***") norm_stats_path = hf_hub_download(repo_id='declare-lab/nora', filename="norm_stats.json") with open(norm_stats_path, "r") as f: self.norm_stats = json.load(f) libero_stats = hf_hub_download(repo_id='moojink/openvla-7b-oft-finetuned-libero-spatial-object-goal-10', filename="dataset_statistics.json") with open(libero_stats, "r") as f: self.norm_stats.update(json.load(f)) def sample_noise(self, shape, device,dtype=torch.float32): noise = torch.normal( mean=0.0, std=1.0, size=shape, dtype=dtype, device=device, ) return noise def sample_time(self, bsize, device,dtype=torch.float32): beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0) time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=dtype) time = time_beta * 0.999 + 0.001 return time def _replace_action_expert_attention(self): """ Iterate through the model's layers and replace the default Qwen2_5_VLAttention with our custom Qwen2_5_VLMoTAttention. """ for i, layer in enumerate(self.action_expert.layers): layer.self_attn = Qwen2_5_VLMoTAttention( config=self.action_expert.config, layer_idx=i ).to(self.action_expert.dtype) layer.self_attn.to(self.action_expert.device) def denoise_step( self, x_t: torch.Tensor, timestep: torch.Tensor, states, vlm_kv_cache: tuple, full_2d_attn_mask: torch.Tensor): """ Applies one denoising step to the noisy action `x_t` at a given `timestep`, conditioned on the VLM's output cache. This function is derived from the main `forward` pass, encapsulating the logic for a single step in the diffusion sampling process. Args: self: The instance of the model class. x_t (torch.Tensor): The noisy action tensor from the previous step. Shape: (batch_size, action_chunk_length, action_dim). timestep (torch.Tensor): The current timestep for each sample in the batch. Shape: (batch_size,). vlm_kv_cache (tuple): The pre-computed key-value cache from the VLM, used as conditioning. vlm_pad_mask (torch.Tensor): The padding mask for the VLM inputs, required to build the cross-attention mask. Shape: (batch_size, vlm_seq_len). Returns: torch.Tensor: The predicted noise `u_t` (epsilon). Shape: (batch_size, action_chunk_length, action_dim). """ device = x_t.device bsz = x_t.shape[0] # 1. Embed the noisy action `x_t` x_t = x_t.to(dtype=self.vlm.dtype) action_input_embeds = self.action_in_proj(x_t) # 2. Create sinusoidal time embeddings from the current timestep time_emb = create_sinusoidal_pos_embedding( timestep, self.lm_expert_config.hidden_size, 4e-3, # Values from your forward pass 4.0, device=device, ) time_emb = time_emb.type(dtype=x_t.dtype) # Expand time embedding to match the action embedding dimensions time_emb = time_emb[:, None, :].expand_as(action_input_embeds) # 3. Combine action and time embeddings and process through MLPs action_time_emb = torch.cat([action_input_embeds, time_emb], dim=2) action_time_emb = self.action_time_mlp_in(action_time_emb) action_time_emb = F.silu(action_time_emb) # swish activation action_time_emb = self.action_time_mlp_out(action_time_emb) if states is not None: states_embed = self.state_emb(states) # print(states_embed.shape,action_input_embeds.shape) states_embed = states_embed.unsqueeze(1).expand_as(action_input_embeds) action_time_emb += states_embed # 4. Construct the attention mask for the action expert. # The expert needs to attend to the VLM context and its own action inputs. # The expert's queries originate from the action sequence, so we slice the mask accordingly. # It can attend to the full VLM context and the action sequence. expert_attention_mask = full_2d_attn_mask[:, -self.action_chunk_length:, :] # 5. Prepare position_ids for the expert. # Note: This implementation mirrors your forward pass, where position_ids for the # expert restart from 0. position_ids = torch.arange(self.action_chunk_length, device=device) # 6. Call the action expert with the prepared inputs and VLM cache. expert_output = self.action_expert( inputs_embeds=action_time_emb, expert_attention_mask=expert_attention_mask.unsqueeze(1).bool(), # Add head dim position_ids=position_ids, vlm_key_values=vlm_kv_cache, use_cache=True, # As in the original forward pass ) # 7. Project the expert's output to get the final noise prediction. velocity = self.action_out_proj(expert_output.last_hidden_state) return velocity def sample_fast_tokens(self,image,image2=None,instruction=None,states=None,unnormalize=False,do_sample=False): device = self.vlm.device states = states.to(device) #states = #print(type(image)) image = resize_image(image) ## IMPORTANT. ENSURE IMAGE RESIZING METHOD IS CONSISTENT WITH PRETRAINIGN #if not isinstance(image, PIL.Image.Image): # image = PIL.Image.fromarray(image) # Construct messages in the expected chat format. Note that nora expects image of size 224 by 224 #image = resize_image(image) if image2 is not None: image2 = resize_image(image2) #if not isinstance(image, PIL.Image.Image): #image = PIL.Image.fromarray(image) # Construct messages in the expected chat format. Note that nora expects image of size 224 by 224 messages = [ { "role": "user", "content": [ { "type": "image", "image": image, "resized_height": 224, "resized_width": 224, },{ "type": "image", "image": image2, "resized_height": 224, "resized_width": 224, }, {"type": "text", "text": instruction}, ], } ] else: messages = [ { "role": "user", "content": [ { "type": "image", "image": image, "resized_height": 224, "resized_width": 224, } , {"type": "text", "text": instruction}, ], } ] # Apply chat template to get the text input for the model text = self.processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) # Process vision information (depends on your process_vision_info function) image_inputs, video_inputs = process_vision_info(messages) # Prepare inputs for the model using the main processor #image_inputs, video_inputs = process_vision_info(messages) inputs = self.processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ) # Move inputs to GPU inputs = {k: v.to(device) for k, v in inputs.items()} generated_ids = self.vlm.generate(**inputs,do_sample=True,temperature=1.0) # --- Extract and Decode Action --- # Find the indices of tokens within the action token range start_idx = (self._ACTION_TOKEN_MIN <= generated_ids[0]) & (generated_ids[0] <= self._ACTION_TOKEN_MAX) start_idx = torch.where(start_idx)[0] if len(start_idx) > 0: start_index = start_idx[0].item() else: start_index = None # or -1 to indicate not found # Extract the first action token ID # Decode the action token using the fast tokenizer # The token ID needs to be map back to the range expected by the fast tokenizer decoder output_action = self.fast_tokenizer.decode([generated_ids[0][start_idx] - self._ACTION_TOKEN_MIN]) return output_action @torch.no_grad() def sample_actions(self, image,image2=None,instruction=None,num_steps:int = 25,states=None,unnorm_key='libero_10',unnormalize=True): """ Generates actions by running the full diffusion sampling process. This function first computes the VLM's key-value cache to use as a conditioning context. It then uses an iterative Euler-method-based sampler, calling `denoise_step` at each timestep to refine a noise tensor into a final action. Args: self: The instance of the model class. vlm_inputs (dict): A dictionary containing the inputs for the VLM, e.g., {'input_ids': ..., 'attention_mask': ...}. noise (Tensor, optional): An initial noise tensor to start the sampling from. If None, it will be sampled randomly. Defaults to None. Shape: (batch_size, action_chunk_length, action_dim). Returns: Tensor: The final, denoised action tensor. Shape: (batch_size, action_chunk_length, action_dim). """ #vlm_inputs = self.prepare_inputs_for_generation(image,instruction) device = self.vlm.device states = states.to(device) #states = #print(type(image)) image = resize_image(image) ## IMPORTANT. ENSURE IMAGE RESIZING METHOD IS CONSISTENT WITH PRETRAINIGN #if not isinstance(image, PIL.Image.Image): # image = PIL.Image.fromarray(image) # Construct messages in the expected chat format. Note that nora expects image of size 224 by 224 #image = resize_image(image) if image2 is not None: image2 = resize_image(image2) #if not isinstance(image, PIL.Image.Image): #image = PIL.Image.fromarray(image) # Construct messages in the expected chat format. Note that nora expects image of size 224 by 224 messages = [ { "role": "user", "content": [ { "type": "image", "image": image, "resized_height": 224, "resized_width": 224, },{ "type": "image", "image": image2, "resized_height": 224, "resized_width": 224, }, {"type": "text", "text": instruction}, ], } ] else: messages = [ { "role": "user", "content": [ { "type": "image", "image": image, "resized_height": 224, "resized_width": 224, } , {"type": "text", "text": instruction}, ], } ] # Apply chat template to get the text input for the model text = self.processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) # Process vision information (depends on your process_vision_info function) image_inputs, video_inputs = process_vision_info(messages) # Prepare inputs for the model using the main processor #image_inputs, video_inputs = process_vision_info(messages) inputs = self.processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ) # Move inputs to GPU inputs = {k: v.to(device) for k, v in inputs.items()} bsz = inputs['input_ids'].shape[0] # 1. Pre-compute the VLM cache. This context is the conditioning for the # entire denoising process and only needs to be computed once. if self.vlm_kv_cache is None: vlm_outputs = self.vlm(**inputs) vlm_kv_cache = vlm_outputs.past_key_values self.vlm_kv_cache = vlm_kv_cache # The VLM's attention mask is its padding mask for the expert. vlm_pad_mask = inputs['attention_mask'].clone() # 2. Initialize the noisy action tensor `x_t`. actions_shape = (bsz, self.action_chunk_length, self.config['max_action_dim']) x_t = self.sample_noise(actions_shape, device=device,dtype=self.vlm.dtype) # 3. Set up the time steps for the Euler solver. # We will step from t=1 down to t=0. #num_steps = self.config.num_steps dt = -1.0 / num_steps dt_tensor = torch.tensor(dt, dtype=self.vlm.dtype, device=device) time = torch.tensor(1.0, dtype=self.vlm.dtype, device=device) states = states.to(self.vlm.dtype) # 4. Iteratively denoise using the Euler method. # The loop continues as long as time is greater than or equal to zero. action_pad_mask = torch.ones(bsz, self.action_chunk_length, device=device).bool() # An all-zero attention mask for the action part allows for full bidirectional attention # within the action chunk, as seen in the original forward pass. action_attn_mask = torch.zeros(bsz, self.action_chunk_length, device=device).bool() # Concatenate VLM (prefix) and action masks. # The VLM's attention mask is its padding mask. concat_pad_mask = torch.cat([vlm_pad_mask, action_pad_mask], dim=1) concat_attn_mask = torch.cat([vlm_pad_mask, action_attn_mask], dim=1) # Create the full 2D attention mask for the combined sequence. full_2d_attn_mask = make_att_2d_masks(concat_pad_mask, concat_attn_mask) while time >= -dt / 2: # Loop until t=0 with torch.no_grad(): # Expand the current time to match the batch size. expanded_time = time.expand(bsz) # Call the denoise_step function to predict the velocity v_t (or noise u_t). # The function takes the current noisy action, timestep, and the # pre-computed VLM cache and padding mask as input. #print(expanded_time) v_t = self.denoise_step( x_t=x_t, timestep=expanded_time, states=states, vlm_kv_cache=self.vlm_kv_cache, full_2d_attn_mask=full_2d_attn_mask, ) # 5. Apply the Euler integration step to update the action tensor. # This moves the action slightly along the direction of the predicted velocity. x_t += dt * v_t time += dt # 6. Return the final denoised action. normalized_action = x_t.cpu().float().numpy() #self.vlm_kv_cache = None if unnormalize is False: return normalized_action action_stats = self._get_action_stats(unnorm_key) mask = action_stats.get("mask", np.ones_like(action_stats["q01"], dtype=bool)) action_high, action_low = np.array(action_stats["q99"]), np.array(action_stats["q01"]) actions = np.where( mask, 0.5 * (normalized_action + 1) * (action_high - action_low) + action_low, normalized_action, ) return actions def _get_action_stats(self, unnorm_key: str) -> Dict[str, Any]: if unnorm_key not in self.norm_stats: raise KeyError( f"The `unnorm_key` '{unnorm_key}' is not in the set of available dataset statistics. " f"Please choose from: {list(self.norm_stats.keys())}" ) return self.norm_stats[unnorm_key]["action"] def forward(self,vlm_inputs, actions,alpha=10.0,use_state=False,states=None ,**kwargs): """ The main forward pass that uses the student model with the expert's cache. """ # The magic happens here: we pass the expert cache into the student's forward call. # This will require modifying how arguments are passed down. ## Precompute the VLM cache with only VLM inputs/attention mask ## Let the Qwen2_5 vlm settle its own attention mask. device = self.vlm.device vlm_outputs = self.vlm( **vlm_inputs, use_cache=True ) vlm_kv_cache = vlm_outputs.past_key_values ## Construct attention mask for the action expert. ## The action expert should be able to attend to the VLM inputs and its own action inputs. ( Prefix + bidirectional attention) bsz = vlm_inputs['input_ids'].shape[0] vlm_pad_mask = vlm_inputs['expert_attention'].clone() vlm_attn_mask = vlm_inputs['attention_mask'].clone() actions = actions.to(self.vlm.dtype) noise = self.sample_noise(actions.shape, actions.device,dtype=actions.dtype) time = self.sample_time(actions.shape[0], actions.device,dtype=actions.dtype) time_expanded = time[:, None, None] x_t = time_expanded * noise + (1 - time_expanded) * actions u_t = noise - actions #x_t = x_t.to(self.vlm.dtype) action_input_embeds = self.action_in_proj(x_t) ## Embed noisy action time_emb = create_sinusoidal_pos_embedding( time, self.lm_expert_config.hidden_size, 4e-3, 4.0, device=device, ) time_emb = time_emb.type(dtype=actions.dtype) time_emb = time_emb[:, None, :].expand_as(action_input_embeds) action_time_emb = torch.cat([action_input_embeds, time_emb], dim=2) ## concat on the hidden size dim action_time_emb = self.action_time_mlp_in(action_time_emb) ## simple linear layer to project back to hidden size dim action_time_emb = F.silu(action_time_emb) # swish == silu action_time_emb = self.action_time_mlp_out(action_time_emb) ## if use_state: states_embed = self.state_emb(states) states_embed = states_embed.unsqueeze(1).expand_as(action_input_embeds) action_time_emb += states_embed action_pad_mask = torch.ones(bsz,self.action_chunk_length,device=device).bool() action_attn_mask = torch.zeros(bsz,self.action_chunk_length,device=device).bool() concat_action_mask = torch.cat([vlm_pad_mask,action_pad_mask],dim=1) concat_attn_mask = torch.cat([vlm_attn_mask,action_attn_mask],dim=1) attn = make_att_2d_masks(concat_action_mask,concat_attn_mask) expert_attention_mask = attn[:, -self.action_chunk_length:, :] position_ids = torch.arange(self.action_chunk_length,device=device) expert_output = self.action_expert(inputs_embeds=action_time_emb, expert_attention_mask=expert_attention_mask.unsqueeze(1).bool(), position_ids= position_ids, vlm_key_values=vlm_kv_cache, use_cache=True) action_out = self.action_out_proj(expert_output.last_hidden_state) expert_loss = alpha*F.mse_loss(action_out, u_t, reduction='mean') loss = expert_loss+ vlm_outputs.loss return {'expert_loss': expert_loss,'combined_loss':loss,'vlm_loss':vlm_outputs.loss}