| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import inspect |
|
|
| import torch |
| from torch import Tensor, nn |
| import torch.nn.functional as F |
|
|
| from starVLA.model.modules.vlm.openpi_transformers.gemma.configuration_gemma import GemmaConfig |
| from starVLA.model.modules.vlm.openpi_transformers.gemma.modeling_gemma import GemmaForCausalLM |
|
|
|
|
| @dataclass |
| class GemmaDims: |
| width: int |
| depth: int |
| mlp_dim: int |
| num_heads: int |
| num_kv_heads: int |
| head_dim: int |
|
|
|
|
| def create_sinusoidal_pos_embedding( |
| time: torch.Tensor, |
| dimension: int, |
| min_period: float, |
| max_period: float, |
| device: torch.device, |
| ) -> Tensor: |
| if dimension % 2 != 0: |
| raise ValueError(f"dimension ({dimension}) must be divisible by 2") |
| if time.ndim != 1: |
| raise ValueError("time must have shape [batch_size]") |
|
|
| fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float32, device=device) |
| period = min_period * (max_period / min_period) ** fraction |
| scaling_factor = 1.0 / period * 2 * torch.pi |
| sin_input = scaling_factor[None, :] * time.to(torch.float32)[:, None] |
| return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) |
|
|
|
|
| def sample_beta(alpha: float, beta: float, bsize: int, device: torch.device) -> torch.Tensor: |
| alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device) |
| beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device) |
| return torch.distributions.Beta(alpha_t, beta_t).sample((bsize,)) |
|
|
|
|
| def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: |
| if att_masks.ndim != 2 or pad_masks.ndim != 2: |
| raise ValueError(f"Expected 2D masks, got {pad_masks.ndim=} {att_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] |
| return att_2d_masks & pad_2d_masks |
|
|
|
|
| def make_att_4d_masks(att_2d_masks: torch.Tensor) -> torch.Tensor: |
| return torch.where(att_2d_masks[:, None, :, :], 0.0, -2.3819763e38) |
|
|
|
|
|
|
| class OpenPIGemma(nn.Module): |
| """Gemma transformer wrapper used by OpenPI action heads.""" |
|
|
| def __init__(self, config: GemmaDims, *, use_adarms: bool = False): |
| super().__init__() |
| hf_config = GemmaConfig( |
| head_dim=config.head_dim, |
| hidden_size=config.width, |
| intermediate_size=config.mlp_dim, |
| num_attention_heads=config.num_heads, |
| num_hidden_layers=config.depth, |
| num_key_value_heads=config.num_kv_heads, |
| vocab_size=257152, |
| hidden_activation="gelu_pytorch_tanh", |
| torch_dtype="float32", |
| use_adarms=bool(use_adarms), |
| adarms_cond_dim=config.width if use_adarms else None, |
| ) |
| gemma = GemmaForCausalLM(config=hf_config) |
| gemma.model.embed_tokens = None |
| self.model = gemma.model |
| self.lm_head = gemma.lm_head |
| self.validate_transformers_api(use_adarms=use_adarms) |
| if use_adarms: |
| self._ensure_adarms_norms(config.width) |
|
|
| def _ensure_adarms_norms(self, cond_dim: int) -> None: |
| for layer in self.model.layers: |
| self._ensure_adarms_norm(layer.input_layernorm, cond_dim) |
| self._ensure_adarms_norm(layer.post_attention_layernorm, cond_dim) |
| self._ensure_adarms_norm(self.model.norm, cond_dim) |
|
|
| @staticmethod |
| def _ensure_adarms_norm(norm: nn.Module, cond_dim: int) -> None: |
| if getattr(norm, "dense", None) is None: |
| width = int(norm.weight.shape[0]) |
| norm.cond_dim = int(cond_dim) |
| norm.dense = nn.Linear(cond_dim, width * 3) |
| nn.init.zeros_(norm.dense.weight) |
| if hasattr(norm, "weight"): |
| norm.weight.data.zero_() |
| norm.weight.requires_grad_(False) |
|
|
| def validate_transformers_api(self, *, use_adarms: bool) -> None: |
| if not use_adarms: |
| return |
| norm_sig = inspect.signature(self.model.layers[0].input_layernorm.forward) |
| model_sig = inspect.signature(self.model.forward) |
| if "cond" not in norm_sig.parameters or "adarms_cond" not in model_sig.parameters: |
| raise RuntimeError( |
| "PI05 requires StarVLA vendored OpenPI-compatible Gemma adaRMS support. " |
| "The imported Gemma does not expose GemmaRMSNorm.forward(cond=...) and " |
| "GemmaModel.forward(adarms_cond=...). " |
| f"Current signatures: GemmaRMSNorm.forward{norm_sig}, GemmaModel.forward{model_sig}" |
| ) |
|
|
|
|
| class OpenPIActionHeadBase(nn.Module): |
| """OpenPI action expert + action head. |
| |
| This module owns the action-side Gemma expert, suffix embedding, action |
| projections, and the suffix-only denoise forward used during inference. |
| """ |
|
|
| use_adarms: bool = False |
|
|
| def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int): |
| super().__init__() |
| self.action_dim = int(action_dim) |
| self.action_horizon = int(action_horizon) |
| self.max_state_dim = int(max_state_dim) |
| self.width = int(config.width) |
| self.gemma_expert = OpenPIGemma(config, use_adarms=self.use_adarms) |
| self.action_in_proj = nn.Linear(self.action_dim, self.width) |
| self.action_out_proj = nn.Linear(self.width, self.action_dim) |
|
|
| @property |
| def model(self): |
| return self.gemma_expert.model |
|
|
| @property |
| def layers(self): |
| return self.gemma_expert.model.layers |
|
|
| @property |
| def norm(self): |
| return self.gemma_expert.model.norm |
|
|
| def set_gemma_expert_eager_attention(self): |
| self.gemma_expert.model.config._attn_implementation = "eager" |
|
|
| def supports_cached_denoise(self) -> bool: |
| return True |
|
|
| def forward_gemma_expert( |
| self, |
| *, |
| inputs_embeds: torch.Tensor, |
| attention_mask: torch.Tensor | None, |
| position_ids: torch.LongTensor | None, |
| past_key_values, |
| use_cache: bool | None, |
| adarms_cond: torch.Tensor | None, |
| ): |
| kwargs = { |
| "inputs_embeds": inputs_embeds, |
| "attention_mask": attention_mask, |
| "position_ids": position_ids, |
| "past_key_values": past_key_values, |
| "use_cache": use_cache, |
| } |
| if "adarms_cond" in inspect.signature(self.gemma_expert.model.forward).parameters: |
| kwargs["adarms_cond"] = adarms_cond |
| elif adarms_cond is not None: |
| raise RuntimeError( |
| "PI05 requires StarVLA vendored OpenPI-compatible Gemma adaRMS support, " |
| "but the imported action expert does not expose it." |
| ) |
| return self.gemma_expert.model.forward(**kwargs) |
|
|
| def sample_action_noise(self, shape, device: torch.device): |
| return torch.normal(0.0, 1.0, size=shape, dtype=torch.float32, device=device) |
|
|
| def sample_flow_time(self, bsize: int, device: torch.device): |
| return (sample_beta(1.5, 1.0, bsize, device) * 0.999 + 0.001).to(torch.float32) |
|
|
| def prepare_flow_targets(self, actions: torch.Tensor, noise=None, time=None): |
| actions = actions.to(device=self.action_in_proj.weight.device, dtype=torch.float32) |
| noise = self.sample_action_noise(actions.shape, actions.device) if noise is None else noise.to(actions.device, torch.float32) |
| time = self.sample_flow_time(actions.shape[0], actions.device) if time is None else time.to(actions.device, torch.float32) |
| x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions |
| u_t = noise - actions |
| return x_t, u_t, noise, time |
|
|
| def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor): |
| raise NotImplementedError |
|
|
| def decode_action_velocity(self, action_hidden_states: torch.Tensor) -> torch.Tensor: |
| action_hidden_states = action_hidden_states[:, -self.action_horizon :].to(self.action_out_proj.weight.dtype) |
| return self.action_out_proj(action_hidden_states).to(torch.float32) |
|
|
| def forward( |
| self, |
| *, |
| attention_mask: torch.Tensor, |
| position_ids: torch.LongTensor, |
| past_key_values, |
| suffix_embs: torch.Tensor, |
| adarms_cond: torch.Tensor | None, |
| force_bfloat16: bool = False, |
| ): |
| if force_bfloat16: |
| suffix_embs = suffix_embs.to(torch.bfloat16) |
| self.set_gemma_expert_eager_attention() |
| suffix_output = self.forward_gemma_expert( |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=suffix_embs, |
| use_cache=False, |
| adarms_cond=adarms_cond, |
| ) |
| suffix_out = suffix_output.last_hidden_state |
| return self.decode_action_velocity(suffix_out) |
|
|
|
|
|
|
| class OpenPI0ActionHead(OpenPIActionHeadBase): |
| """PI0 suffix path: state token plus action-time MLP, no adaRMS conditioning.""" |
|
|
| use_adarms = False |
|
|
| def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int): |
| super().__init__(config, action_dim=action_dim, action_horizon=action_horizon, max_state_dim=max_state_dim) |
| self.state_proj = nn.Linear(self.max_state_dim, self.width) |
| self.action_time_mlp_in = nn.Linear(2 * self.width, self.width) |
| self.action_time_mlp_out = nn.Linear(self.width, self.width) |
|
|
| def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor): |
| state_emb = self.state_proj(state.to(self.state_proj.weight.dtype)) |
| embs = [state_emb[:, None, :]] |
| pad_masks = [torch.ones(state_emb.shape[0], 1, dtype=torch.bool, device=state_emb.device)] |
| att_masks = [1] |
|
|
| time_emb = create_sinusoidal_pos_embedding( |
| timestep, |
| self.width, |
| min_period=4e-3, |
| max_period=4.0, |
| device=timestep.device, |
| ).to(self.action_in_proj.weight.dtype) |
| action_emb = self.action_in_proj(noisy_actions.to(self.action_in_proj.weight.dtype)) |
| time_emb = time_emb[:, None, :].expand_as(action_emb) |
| action_time_emb = self.action_time_mlp_out( |
| F.silu(self.action_time_mlp_in(torch.cat([action_emb, time_emb], dim=2))) |
| ) |
| embs.append(action_time_emb) |
| pad_masks.append(torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=action_time_emb.device)) |
| att_masks += [1] + ([0] * (self.action_horizon - 1)) |
| embs = torch.cat(embs, dim=1) |
| pad_masks = torch.cat(pad_masks, dim=1) |
| att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device)[None, :] |
| return embs, pad_masks, att_masks.expand(embs.shape[0], -1), None |
|
|
|
|
| class OpenPI05ActionHead(OpenPIActionHeadBase): |
| """PI05 suffix path: action tokens conditioned by adaRMS time embedding.""" |
|
|
| use_adarms = True |
|
|
| def __init__(self, config: GemmaDims, *, action_dim: int, action_horizon: int, max_state_dim: int): |
| super().__init__(config, action_dim=action_dim, action_horizon=action_horizon, max_state_dim=max_state_dim) |
| self.time_mlp_in = nn.Linear(self.width, self.width) |
| self.time_mlp_out = nn.Linear(self.width, self.width) |
|
|
| def embed_action_suffix(self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor): |
| time_emb = create_sinusoidal_pos_embedding( |
| timestep, |
| self.width, |
| min_period=4e-3, |
| max_period=4.0, |
| device=timestep.device, |
| ).to(self.time_mlp_in.weight.dtype) |
| action_time_emb = self.action_in_proj(noisy_actions.to(self.action_in_proj.weight.dtype)) |
| adarms_cond = F.silu(self.time_mlp_out(F.silu(self.time_mlp_in(time_emb)))) |
| pad_masks = torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=action_time_emb.device) |
| att_masks = torch.tensor([1] + ([0] * (self.action_horizon - 1)), dtype=torch.bool, device=action_time_emb.device) |
| return action_time_emb, pad_masks, att_masks[None, :].expand(action_time_emb.shape[0], -1), adarms_cond |
|
|
