""" Qwen-OFT Multi-Robot Framework Extends Qwen-OFT to support mixed-benchmark co-training where different robots may have different action/state dimensions. The framework pads action/state to a unified max dimension and injects robot meta information into instruction. """ from typing import List, Optional import torch import torch.nn as nn import numpy as np from starVLA.training.trainer_utils import initialize_overwatch from starVLA.model.tools import FRAMEWORK_REGISTRY, collate_fn_extend_dim from deployment.model_server.tools.image_tools import to_pil_preserve from starVLA.model.framework.base_framework import baseframework from starVLA.model.modules.vlm import get_vlm_model from starVLA.model.modules.action_model.MLP_ActionHeader import get_action_model from starVLA.training.trainer_utils.trainer_tools import resize_images logger = initialize_overwatch(__name__) IGNORE_INDEX = -100 # TODO make this configurable from yaml multi_robot_action_heads = { "franka": {"action_dim": 7, "NUM_ACTIONS_CHUNK": 16, "robo_info": "single arm, delta eef", "robot_mlp_id": 1}, "oxe_bridge": {"action_dim": 7, "NUM_ACTIONS_CHUNK": 16, "robo_info": "single arm, delta eef", "robot_mlp_id": 2}, "oxe_rt1": {"action_dim": 7, "NUM_ACTIONS_CHUNK": 16, "robo_info": "single arm, delta eef", "robot_mlp_id": 3}, "gr1": {"action_dim": 29, "NUM_ACTIONS_CHUNK": 16, "robo_info": "dual-arm dexterous hands, Joint", "robot_mlp_id": 4}, "robotwin": {"action_dim": 14, "NUM_ACTIONS_CHUNK": 16, "robo_info": "arm gripper, Joint", "robot_mlp_id": 5}, } @FRAMEWORK_REGISTRY.register("QwenOFT_xrobot") class Qwenvl_OFT_multiRobo(baseframework): def __init__( self, config: Optional[dict] = None, **kwargs, ) -> None: super().__init__() self.config = config self.qwen_vl_interface = get_vlm_model(config=self.config) # Align action head hidden dim with the VLM hidden size. config.framework.action_model.action_hidden_dim = self.qwen_vl_interface.model.config.hidden_size self.action_model = get_action_model(config=self.config) self.future_action_window_size = config.framework.action_model.future_action_window_size self.past_action_window_size = config.framework.action_model.past_action_window_size self.chunk_len = self.past_action_window_size + 1 + self.future_action_window_size self.action_token = "🔍" self.action_token_id = self.qwen_vl_interface.processor.tokenizer( self.action_token, add_special_tokens=False )["input_ids"][0] self.l1_loss = nn.L1Loss() def forward( self, examples: List[dict] = None, **kwargs, ) -> dict: if type(examples) is not list: examples = [examples] # Pad action/state dims for mixed-benchmark co-training. examples = collate_fn_extend_dim(examples, max_dim=self.config.framework.action_model.action_dim) batch_images = [example["image"] for example in examples] instructions = [example["lang"] for example in examples] actions = [example["action"] for example in examples] instructions = self._add_robo_meta_tokens_to_instructions(examples, instructions) action_tokens = self.action_token * self.chunk_len prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: {action_tokens}." instructions = [instruction + prompt_suffix for instruction in instructions] qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) with torch.autocast("cuda", dtype=torch.bfloat16): qwenvl_outputs = self.qwen_vl_interface( **qwen_inputs, output_attentions=False, output_hidden_states=True, return_dict=True, ) last_hidden = qwenvl_outputs.hidden_states[-1] with torch.autocast("cuda", dtype=torch.float32): input_ids = qwen_inputs.get("input_ids", None) action_queries = self._gather_action_token_embeddings( last_hidden, input_ids, action_token_id=self.action_token_id ) pred_actions = self.action_model.predict_action(action_queries) actions = torch.tensor(np.array(actions), device=pred_actions.device, dtype=pred_actions.dtype) actions_target = actions[:, -(self.future_action_window_size + 1) :, :] action_loss = self.l1_loss(pred_actions, actions_target) return {"action_loss": action_loss} @torch.inference_mode() def predict_action( self, examples: List[dict] = None, **kwargs: str, ) -> dict: if type(examples) is not list: examples = [examples] # Keep inference input format consistent with training. examples = collate_fn_extend_dim(examples, max_dim=self.config.framework.action_model.action_dim) batch_images = [to_pil_preserve(example["image"]) for example in examples] instructions = [example["lang"] for example in examples] instructions = self._add_robo_meta_tokens_to_instructions(examples, instructions) train_obs_image_size = getattr(self.config.datasets.vla_data, "image_size", None) if train_obs_image_size: batch_images = resize_images(batch_images, target_size=train_obs_image_size) action_tokens = self.action_token * self.chunk_len prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: {action_tokens}." instructions = [instruction + prompt_suffix for instruction in instructions] qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) with torch.autocast("cuda", dtype=torch.bfloat16): qwenvl_outputs = self.qwen_vl_interface( **qwen_inputs, output_attentions=False, output_hidden_states=True, return_dict=True, ) last_hidden = qwenvl_outputs.hidden_states[-1] with torch.autocast("cuda", dtype=torch.float32): input_ids = qwen_inputs.get("input_ids", None) action_queries = self._gather_action_token_embeddings( last_hidden, input_ids, action_token_id=self.action_token_id ) pred_actions = self.action_model.predict_action(action_queries) normalized_actions = pred_actions.detach().cpu().numpy() return {"normalized_actions": normalized_actions} def _add_robo_meta_tokens_to_instructions(self, examples: List[dict], instructions: List[str]) -> List[str]: enhanced_instructions = [] for example, instruction in zip(examples, instructions): robot_tag = example.get("robot_tag", "franka") if robot_tag not in multi_robot_action_heads: raise ValueError( f"Unknown robot_tag `{robot_tag}`. Available tags: {list(multi_robot_action_heads.keys())}" ) robot_meta = multi_robot_action_heads[robot_tag] robo_info = "Robot name: {}. Action Dim: {}. Robot info: {}".format( robot_tag, robot_meta["action_dim"], robot_meta.get("robo_info", "N/A"), ) chunk_len = robot_meta["NUM_ACTIONS_CHUNK"] prompt_suffix = f" Please predict the next {chunk_len} robot actions for the robot {robo_info}." enhanced_instructions.append(instruction + prompt_suffix) return enhanced_instructions def _gather_action_token_embeddings( self, last_hidden: torch.Tensor, input_ids: torch.Tensor, action_token_id=None, ) -> torch.Tensor: if action_token_id is None: raise ValueError("action_token_id must not be None") device = input_ids.device batch_size, seq_len, hidden_dim = last_hidden.shape if isinstance(action_token_id, (list, tuple, set)): id_list = torch.tensor(list(action_token_id), device=device, dtype=input_ids.dtype) mask = torch.isin(input_ids, id_list) else: mask = input_ids == action_token_id counts = mask.sum(dim=1) if (counts < self.chunk_len).any(): insufficient = (counts < self.chunk_len).nonzero(as_tuple=False).flatten().tolist() raise RuntimeError( f"Some samples have fewer than {self.chunk_len} action tokens: {insufficient} | counts={counts.tolist()}" ) idx = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, seq_len) masked_pos = torch.where(mask, idx, torch.full_like(idx, -1)) topk_pos = masked_pos.topk(k=self.chunk_len, dim=-1).values selected_pos = topk_pos.sort(dim=-1).values expanded_index = selected_pos.unsqueeze(-1).expand(-1, -1, hidden_dim) action_queries = last_hidden.gather(dim=1, index=expanded_index) return action_queries