import os from scipy.fft import dct from scipy.fft import idct import yaml import torch import numpy as np import dataclasses import copy import json from PIL import Image from safetensors.torch import load_file from qwen_vl_utils.vision_process import smart_resize from transformers import BatchFeature, AutoProcessor from wall_x.model.action_head import Normalizer from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof from numba import jit, prange try: from spatial_tokenizer.spatial_tokenizer_kdisk import SpatialActionTokenizer except ImportError: SpatialActionTokenizer = None device = "cuda" dof_config = { "follow_left_ee_cartesian_pos": 3, "follow_left_ee_rotation": 3, "follow_left_gripper": 1, "follow_right_ee_cartesian_pos": 3, "follow_right_ee_rotation": 3, "follow_right_gripper": 1, "head_actions": 2, "height": 1, "car_pose": 3, "velocity_decomposed": 3, } _CAM_NAME_MAPPING = { "face_view": "front view", "left_wrist_view": "left wrist view", "right_wrist_view": "right wrist view", "move1_view": "move view", "move2_view": "move view", "wall_view": "wall view", "top_view": "top view", "side_view": "side view", "global_view": "global view", } camera_to_view_mapping = { "camera_front": "face_view", "camera_left": "left_wrist_view", "camera_right": "right_wrist_view", "camera_side": "side_view", "camera_global": "global_view", } action_key_mapping = { "follow_left_ee_cartesian_pos": "follow1_pos[:3]", "follow_left_ee_rotation": "follow1_pos[3:6]", "follow_left_gripper": "follow1_pos[6:7]", "follow_right_ee_cartesian_pos": "follow2_pos[:3]", "follow_right_ee_rotation": "follow2_pos[3:6]", "follow_right_gripper": "follow2_pos[6:7]", "head_actions": "head_pos", "height": "lift", "velocity_decomposed": "velocity_decomposed", "follow_left_arm_joint_cur": "follow1_joints_cur[-1:]", "follow_right_arm_joint_cur": "follow2_joints_cur[-1:]", "follow_left_arm_joint_pos": "follow1_pos", "follow_right_arm_joint_pos": "follow2_pos", } dim_dof_config = { "right_xyz": {"rpy": (0, 3), "so3": (0, 3)}, "right_rot": {"rpy": (3, 6), "so3": (3, 9)}, "right_gripper": {"rpy": (6, 7), "so3": (9, 10)}, } SINGLE_ARM_DIM = 7 @jit(nopython=True, parallel=True) def euler_to_matrix_zyx_batch_nb(eulers): N = eulers.shape[0] R = np.empty((N, 3, 3), dtype=np.float64) for i in prange(N): roll = eulers[i, 0] pitch = eulers[i, 1] yaw = eulers[i, 2] cy, sy = np.cos(yaw), np.sin(yaw) cp, sp = np.cos(pitch), np.sin(pitch) cr, sr = np.cos(roll), np.sin(roll) R[i, 0, 0] = cy * cp R[i, 0, 1] = cy * sp * sr - sy * cr R[i, 0, 2] = cy * sp * cr + sy * sr R[i, 1, 0] = sy * cp R[i, 1, 1] = sy * sp * sr + cy * cr R[i, 1, 2] = sy * sp * cr - cy * sr R[i, 2, 0] = -sp R[i, 2, 1] = cp * sr R[i, 2, 2] = cp * cr return R @jit(nopython=True, parallel=True) def compose_state_and_delta_to_abs_rpy(delta, state): """ Input: delta: (N,3) -> Δrpy(ZYX) or (N,6) -> Δ6D (first two rows flattened) state: (3,) -> rpy(ZYX) or (6,) -> 6D (first two rows flattened) Output: abs_rpy: (N,3) Absolute pose rpy(ZYX, radians), normalized to (-π, π] """ if delta.shape[-1] == 3: R_delta = euler_to_matrix_zyx_batch_nb(delta) # (N,3,3) elif delta.shape[-1] == 6: R_delta = so3_to_matrix_batch_nb(delta) # (N,3,3) else: raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}") if state.shape[-1] == 3: R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] # (3,3) elif state.shape[-1] == 6: R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] # (3,3) else: raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}") N = R_delta.shape[0] R_abs = np.empty((N, 3, 3), dtype=np.float64) S00 = R_state[0, 0] S01 = R_state[0, 1] S02 = R_state[0, 2] S10 = R_state[1, 0] S11 = R_state[1, 1] S12 = R_state[1, 2] S20 = R_state[2, 0] S21 = R_state[2, 1] S22 = R_state[2, 2] for i in prange(N): A00 = R_delta[i, 0, 0] A01 = R_delta[i, 0, 1] A02 = R_delta[i, 0, 2] A10 = R_delta[i, 1, 0] A11 = R_delta[i, 1, 1] A12 = R_delta[i, 1, 2] A20 = R_delta[i, 2, 0] A21 = R_delta[i, 2, 1] A22 = R_delta[i, 2, 2] R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 @jit(nopython=True, parallel=True) def so3_to_matrix_batch_nb(batch_so3): N = batch_so3.shape[0] R_all = np.empty((N, 3, 3), dtype=np.float64) eps = 1e-12 for i in prange(N): r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2] r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5] # normalize r1 n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps r1x /= n1 r1y /= n1 r1z /= n1 # orthogonalize r2 to r1, then normalize dot12 = r1x * r2x + r1y * r2y + r1z * r2z r2x -= dot12 * r1x r2y -= dot12 * r1y r2z -= dot12 * r1z n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps r2x /= n2 r2y /= n2 r2z /= n2 # r3 = r1 x r2 r3x = r1y * r2z - r1z * r2y r3y = r1z * r2x - r1x * r2z r3z = r1x * r2y - r1y * r2x R_all[i, 0, 0] = r1x R_all[i, 0, 1] = r1y R_all[i, 0, 2] = r1z R_all[i, 1, 0] = r2x R_all[i, 1, 1] = r2y R_all[i, 1, 2] = r2z R_all[i, 2, 0] = r3x R_all[i, 2, 1] = r3y R_all[i, 2, 2] = r3z return R_all @jit(nopython=True, parallel=True) def matrix_to_euler_zyx_batch_nb(Rs): """ R = Rz(yaw) * Ry(pitch) * Rx(roll) extract: pitch = asin(-R[2,0]) roll = atan2(R[2,1], R[2,2]) yaw = atan2(R[1,0], R[0,0]) """ N = Rs.shape[0] eulers = np.empty((N, 3), dtype=np.float64) for i in prange(N): r00 = Rs[i, 0, 0] # r01 = Rs[i, 0, 1] # r02 = Rs[i, 0, 2] r10 = Rs[i, 1, 0] # r11 = Rs[i, 1, 1] # r12 = Rs[i, 1, 2] r20 = Rs[i, 2, 0] r21 = Rs[i, 2, 1] r22 = Rs[i, 2, 2] x = -r20 if x > 1.0: x = 1.0 elif x < -1.0: x = -1.0 pitch = np.arcsin(x) roll = np.arctan2(r21, r22) yaw = np.arctan2(r10, r00) eulers[i, 0] = roll eulers[i, 1] = pitch eulers[i, 2] = yaw return eulers @jit(nopython=True, parallel=True) def canonicalize_euler_zyx_batch_nb(rpy_batch): """ Batch ZYX Euler Angle Normalization (Parallel Version) Input: rpy_batch (N, 3) [roll, pitch, yaw] (radians) Output: out (N, 3) Constrained to the same branch with each component in (-π, π] Rules: 1) First, wrap each component to (-π, π] 2) If p > π/2: p = π - p; r += π; y += π If p <= -π/2: p = -π - p; r += π; y += π 3) Finally, wrap each component to (-π, π] again. """ N = rpy_batch.shape[0] out = np.empty_like(rpy_batch) two_pi = 2.0 * np.pi for i in prange(N): r = rpy_batch[i, 0] p = rpy_batch[i, 1] y = rpy_batch[i, 2] r = (r + np.pi) % two_pi - np.pi p = (p + np.pi) % two_pi - np.pi y = (y + np.pi) % two_pi - np.pi if p > np.pi / 2.0: p = np.pi - p r = r + np.pi y = y + np.pi elif p <= -np.pi / 2.0: p = -np.pi - p r = r + np.pi y = y + np.pi r = (r + np.pi) % two_pi - np.pi p = (p + np.pi) % two_pi - np.pi y = (y + np.pi) % two_pi - np.pi out[i, 0] = r out[i, 1] = p out[i, 2] = y return out def so3_to_euler_zyx_batch_nb(batch_so3): matrix = so3_to_matrix_batch_nb(batch_so3) eulers = matrix_to_euler_zyx_batch_nb(matrix) return canonicalize_euler_zyx_batch_nb(eulers) def update_model_config(train_config, model_config): model_config.use_state_string_representation = train_config["data"].get( "use_state_string_representation", False ) model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) model_config.dof_config = train_config["dof_config"] model_config.agent_pos_config = train_config["agent_pos_config"] model_config.action_horizon_flow = train_config["data"].get( "action_horizon_flow", 32 ) if train_config.get("_attn_implementation", None) is not None: model_config._attn_implementation = train_config["_attn_implementation"] return model_config def move_to_cuda(obj, device=device): if isinstance(obj, torch.Tensor): return obj.to(device) elif isinstance(obj, (dict, BatchFeature)): return {k: move_to_cuda(v, device) for k, v in obj.items()} elif isinstance(obj, list): return [move_to_cuda(v, device) for v in obj] elif isinstance(obj, tuple): return tuple(move_to_cuda(v, device) for v in obj) else: return obj def extract_components(data, config, is_rpy): mode = "rpy" if is_rpy else "so3" result = {} for key, slice_config in config.items(): slice_range = slice_config[mode] result[key] = data[:, slice_range[0] : slice_range[1]] return result @dataclasses.dataclass class WallxInferArgs: config_path: str | None = None checkpoint_path: str | None = None action_mode: str = "diffusion" # "ar" or "diffusion" max_time_step: int = 1000 action_start_ratio: float = 0 action_end_ratio: float = 0.6 model_action_dim: int = 14 action_horizon: int = 32 action_dim: int = 14 interpolate_action: bool = False interpolate_multiplier: int = 1 turtle_as_desktop: bool = False generate_subtask: bool = False subtask_interval: int = 0 with_cur: bool = False state_str: bool = True ### NOTE wostate: bool = False delta_action: bool = False state_rpy: bool = True action_rpy: bool = True dataset_name: str = "robochallenge_aloha" use_hard_prompt: bool = True dct_scale: float = -1 class WallxModelWrapper: def __init__(self, args: WallxInferArgs): self.args = args self.get_model_and_processor() self.action_predict_mode = "ar" if args.action_mode == "ar" else "diffusion" print("action_predict_mode", self.action_predict_mode, flush=True) def get_model_and_processor(self): if self.args.config_path is None: self.args.config_path = os.path.join( self.args.checkpoint_path, "config.yml" ) with open(self.args.config_path, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) self.config = config self.dof_config = config["dof_config"] self.agent_pos_config = config["agent_pos_config"] self.obs_action_keys = [ "follow_left_ee_cartesian_pos", "follow_left_ee_rotation", "follow_left_gripper", "follow_right_ee_cartesian_pos", "follow_right_ee_rotation", "follow_right_gripper", ] print("obs_action_keys", self.obs_action_keys, flush=True) self.action_tokenizer_type = config.get("action_tokenizer_type", None) config_path = config["qwen_vl_act_config_path"] self.processor = AutoProcessor.from_pretrained( config["processor_path"], use_fast=True ) self.processor.tokenizer.padding_side = "left" new_tokens = ["<|propri|>", "<|action|>"] # load fast tokenizer self.action_tokenizer_type = config.get("action_tokenizer_type", None) print("self.action_tokenizer_type", self.action_tokenizer_type, flush=True) self.action_tokenizer = None self.action_mapper = None if self.action_tokenizer_type: # fast if self.action_tokenizer_type == "fast": print("Using fast tokenizer") self.action_tokenizer = AutoProcessor.from_pretrained( config["action_tokenizer_path"], trust_remote_code=True ) elif self.action_tokenizer_type == "spatialvla": print("Using spatialvla tokenizer") assert ( SpatialActionTokenizer is not None ), "SpatialActionTokenizer is not installed" self.action_tokenizer = SpatialActionTokenizer() else: raise ValueError( f"Unsupported action tokenizer type: {self.action_tokenizer_type}" ) new_tokens += [ f"<|action_token_{i}|>" for i in range(self.action_tokenizer.vocab_size) ] # num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens) # define action mapper if self.action_tokenizer_type: self.action_mapper = {} for i in range(self.action_tokenizer.vocab_size): token = f"<|action_token_{i}|>" token_id = self.processor.tokenizer.convert_tokens_to_ids(token) self.action_mapper[token_id] = i # action & propri normalizer self._register_normalizers() model_type = config["model_type"] if model_type == "qwen2_5": print("Using qwen2_5 model as base model") from wall_x.model.qwen2_5_based import ( Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig, ) ModelClass = Qwen2_5_VLMoEForAction ConfigClass = Qwen2_5_VLConfig model_config = ConfigClass.from_pretrained(config_path) model_config = update_model_config(config, model_config) print("model_config", model_config, flush=True) # if self.args.action_mode == "ar": # model_config._attn_implementation = "flash_attention_2" # else model_config._attn_implementation = "sdpa" model_config.vision_config._attn_implementation = "flash_attention_2" if model_config.model_type == "qwen3_vl": self.MAX_PIXELS = 16384 * 32 * 32 self.MIN_PIXELS = 4 * 32 * 32 self.IMAGE_FACTOR = 32 elif model_config.model_type == "qwen2_5_vl": self.MAX_PIXELS = 16384 * 28 * 28 self.MIN_PIXELS = 4 * 28 * 28 self.IMAGE_FACTOR = 28 model = ModelClass( model_config, self.action_tokenizer_type, self.processor, self.action_tokenizer, self.action_mapper, ) model.resize_token_embeddings(len(self.processor.tokenizer)) state_dict = load_file( self.args.checkpoint_path + "/model.safetensors", device="cpu" ) if os.path.exists(os.path.join(self.args.checkpoint_path, "global_step.pth")): global_step = torch.load( os.path.join(self.args.checkpoint_path, "global_step.pth") )["global_step"] print("Loaded global step:", global_step) msg = model.load_state_dict(state_dict, strict=False) print(msg) model.eval() model.set_normalizer( copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) ) model.to(device) model.to_bfloat16_for_selected_params() self.model = model print("self.args.dataset_name", self.args.dataset_name, flush=True) print( "normalizer_action min", self.normalizer_action.min.__getattr__(self.args.dataset_name), flush=True, ) print( "normalizer_action delta", self.normalizer_action.delta.__getattr__(self.args.dataset_name), flush=True, ) def _register_normalizers(self): if self.config.get("customized_action_statistic_dof", None): action_statistic_dof = json.load( open(self.config["customized_action_statistic_dof"], "r") ) else: action_statistic_dof = default_action_statistic_dof if os.path.exists(self.args.checkpoint_path + "/normalizer_action.pth"): print( "Loading normalizer_action from checkpoint", self.args.checkpoint_path + "/normalizer_action.pth", flush=True, ) self.normalizer_action = Normalizer.from_ckpt( self.args.checkpoint_path + "/normalizer_action.pth" ) else: self.normalizer_action = Normalizer( action_statistic_dof, self.config["dof_config"], min_key=self.config.get("min_key", "min"), delta_key=self.config.get("delta_key", "delta"), ) # print("action_statistic_dof",action_statistic_dof) if os.path.exists(self.args.checkpoint_path + "/normalizer_propri.pth"): print( "Loading normalizer_propri from checkpoint", self.args.checkpoint_path + "/normalizer_propri.pth", flush=True, ) self.normalizer_propri = Normalizer.from_ckpt( self.args.checkpoint_path + "/normalizer_propri.pth" ) else: self.normalizer_propri = Normalizer( action_statistic_dof, self.config["agent_pos_config"], min_key=self.config.get("min_key", "min"), delta_key=self.config.get("delta_key", "delta"), ) print("self.args.dataset_name", self.args.dataset_name, flush=True) print( "normalizer_propri min", self.normalizer_propri.min.__getattr__(self.args.dataset_name), flush=True, ) print( "normalizer_propri delta", self.normalizer_propri.delta.__getattr__(self.args.dataset_name), flush=True, ) print( "normalizer_action min", self.normalizer_action.min.__getattr__(self.args.dataset_name), flush=True, ) print( "normalizer_action delta", self.normalizer_action.delta.__getattr__(self.args.dataset_name), flush=True, ) def get_text_ar(self, instruction, camera_names, norm_state=None, state_mask=None): role_start_symbol = "<|im_start|>" role_end_symbol = "<|im_end|>" vision_start_symbol = "<|vision_start|>" vision_end_symbol = "<|vision_end|>" image_pad_symbol = "<|image_pad|>" propri_symbol = "<|propri|>" prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" user_request = f"{role_start_symbol}user\nObservation:" print("camera_names", camera_names, flush=True) for cam_name in camera_names: user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" user_request += "\nInstruction:" if self.args.state_str: assert norm_state is not None if isinstance(norm_state, torch.Tensor): if state_mask is not None: if isinstance(state_mask, torch.Tensor): mask_1d = state_mask[0, 0].to( dtype=torch.bool, device=norm_state.device ) else: mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ 0, 0 ].to(dtype=torch.bool) norm_state = norm_state[..., mask_1d] norm_state = norm_state.detach().cpu().numpy() print("norm_state", norm_state, flush=True) discretized_state = ( np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 ) propri = " ".join(map(str, discretized_state[0, 0])) elif self.args.wostate: propri = "" else: propri = propri_symbol text_prompt = ( f"\nPredict the next action in robot action.\nProprioception: {propri}\n" ) user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" assistant_message = f"{role_start_symbol}assistant\n" text = prologue + user_message + assistant_message return text def get_text_flow( self, instruction, camera_names, action_chunk_size, norm_state=None, state_mask=None, ): role_start_symbol = "<|im_start|>" role_end_symbol = "<|im_end|>" vision_start_symbol = "<|vision_start|>" vision_end_symbol = "<|vision_end|>" image_pad_symbol = "<|image_pad|>" propri_symbol = "<|propri|>" action_symbol = "<|action|>" action_space = "Rel EEF" if self.args.delta_action else "Abs EEF" _camera = ", ".join([_CAM_NAME_MAPPING[cam_name] for cam_name in camera_names]) prologue = f"<|im_start|>system\nYou are an embodied vision-language-action (VLA) model controlling the robot with language instructions.\n Embodiment: {self.args.dataset_name.split('_')[-1]}\n Camera Setup: {_camera},\n Frequency: 32HZ\n Action Space: {action_space}\n<|im_end|>\n" user_request = f"{role_start_symbol}user\nObservation:" print("camera_names", camera_names, flush=True) for cam_name in camera_names: user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" user_request += "\nInstruction:" if self.args.state_str: assert norm_state is not None if isinstance(norm_state, torch.Tensor): if state_mask is not None: if isinstance(state_mask, torch.Tensor): mask_1d = state_mask[0, 0].to( dtype=torch.bool, device=norm_state.device ) else: mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ 0, 0 ].to(dtype=torch.bool) norm_state = norm_state[..., mask_1d] norm_state = norm_state.detach().cpu().numpy() discretized_state = ( np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 ) propri = " ".join(map(str, discretized_state[0, 0])) elif self.args.wostate: propri = "" else: propri = propri_symbol text_prompt = ( f"\nPredict the next action in robot action.\nProprioception: {propri}\n" ) user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" assistant_message = f"{role_start_symbol}assistant\n" action = f"{action_symbol * action_chunk_size}" text = prologue + user_message + assistant_message + action return text def resize_images(self, observation): image_inputs = [] view_candidates = [ "face_view", "left_wrist_view", "right_wrist_view", "side_view", "global_view", ] for key in observation.keys(): if key not in view_candidates: print("!!! key not in view_candidates", key, flush=True) continue # 1. Get the original image current_obs = observation[key] img_pil = Image.fromarray(current_obs) orig_width, orig_height = img_pil.size # 2. Apply resolution limits (if the configuration is not -1) target_size = 256 if target_size != -1: # Logic for maintaining aspect ratio constraints if orig_width > orig_height: new_width = target_size new_height = int(target_size * orig_height / orig_width) else: new_height = target_size new_width = int(target_size * orig_width / orig_height) img_pil = img_pil.resize((new_width, new_height)) # 3. Apply intelligent scaling current_width, current_height = img_pil.size resized_height, resized_width = smart_resize( current_height, current_width, factor=self.IMAGE_FACTOR, min_pixels=self.MIN_PIXELS, max_pixels=self.MAX_PIXELS, ) resized_img = img_pil.resize((resized_width, resized_height)) print("resized_img", resized_img.size, flush=True) image_inputs.append(resized_img) return image_inputs def _construct_input( self, observation, instruction, camera_names, valid_action_dim=7, mode="ar", single_image=False, ): additional_inputs = {} agent_pos = torch.from_numpy(observation["agent_pos"]) agent_pos_mask = torch.from_numpy(observation["agent_pos_mask"]) dof_mask = torch.from_numpy(observation["dof_mask"]) additional_inputs["dof_mask"] = dof_mask print("before normalizing agent_pos", agent_pos, flush=True) if self.normalizer_propri is not None: agent_pos = self.normalizer_propri.normalize_data( agent_pos, [self.args.dataset_name] ) additional_inputs["proprioception"] = agent_pos additional_inputs["agent_pos_mask"] = agent_pos_mask print( f"normalizing agent_pos: {agent_pos}, {self.args.dataset_name}", flush=True, ) print("agent_pos_mask", agent_pos_mask, flush=True) print("dof_mask", dof_mask[0, 0], flush=True) if mode == "ar": text = self.get_text_ar( instruction, camera_names, agent_pos, agent_pos_mask ) elif mode == "diffusion": text = self.get_text_flow( instruction, camera_names, self.args.action_horizon, agent_pos, agent_pos_mask, ) elif mode == "subtask": text = self.get_text_subtask(instruction, single_view=single_image) else: raise ValueError(f"Invalid mode: {mode}") text = [text] image_inputs = self.resize_images(observation) if single_image: image_inputs = [ image_inputs[0] ] # single view subtask/vqa use head view only image_inputs = self.processor.image_processor( images=image_inputs, videos=None, return_tensors="pt" ) image_grid_thw = image_inputs["image_grid_thw"] # Processing image placeholder tokens in the text if image_grid_thw is not None: merge_length = self.processor.image_processor.merge_size**2 index = 0 for i in range(len(text)): while "<|image_pad|>" in text[i]: # Replace image placeholders with actual quantities. text[i] = text[i].replace( "<|image_pad|>", "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length), 1, ) index += 1 text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") text_inputs = self.processor.tokenizer( text, return_tensors="pt", padding=True, truncation=True, max_length=1024 ) inputs = BatchFeature(data={**text_inputs, **image_inputs}) action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id additional_inputs["dataset_names"] = [self.args.dataset_name] inputs.update(additional_inputs) inputs = move_to_cuda(inputs, device) print("inputs", inputs.keys(), flush=True) for k in inputs.keys(): if isinstance(inputs[k], torch.Tensor): print(k, inputs[k].shape, flush=True) return inputs def preprocess(self, state, views, valid_action_dim): # print("state",state, flush=True) # state: dict, keys: follow1_pos, follow2_pos model_action_dim = sum(self.dof_config.values()) if valid_action_dim not in (SINGLE_ARM_DIM, 2 * SINGLE_ARM_DIM): raise ValueError( f"Invalid valid_action_dim: {valid_action_dim}, expect 7 or 14" ) # 1) First, prepare a container of size (1, 1, D), where D = model_action_dim. agent_data = np.zeros((1, 1, model_action_dim), dtype=np.float32) agent_pos_mask = np.zeros((1, 1, model_action_dim), dtype=np.float32) dof_mask = np.zeros( (1, self.args.action_horizon, model_action_dim), dtype=np.float32 ) # 2) Determine the interval to be filled [start:end) if valid_action_dim == SINGLE_ARM_DIM: start = 0 if model_action_dim == SINGLE_ARM_DIM else SINGLE_ARM_DIM end = start + SINGLE_ARM_DIM if end > model_action_dim: raise ValueError( f"model_action_dim={model_action_dim} too small for valid_action_dim=7 " f"(need end={end})" ) follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( 1, 1, SINGLE_ARM_DIM ) agent_data[:, :, start:end] = follow2 agent_pos_mask[:, :, start:end] = 1 dof_mask[:, :, start:end] = 1 else: # valid_action_dim == 14 end = 2 * SINGLE_ARM_DIM if end > model_action_dim: raise ValueError( f"model_action_dim={model_action_dim} too small for valid_action_dim=14" ) follow1 = np.asarray(state["follow1_pos"], dtype=np.float32).reshape( 1, 1, SINGLE_ARM_DIM ) follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( 1, 1, SINGLE_ARM_DIM ) agent_data[:, :, :end] = np.concatenate([follow1, follow2], axis=-1) agent_pos_mask[:, :, :end] = 1 dof_mask[:, :, :end] = 1 observation = { camera_to_view_mapping[key]: views[key][0] for key in views.keys() } observation["agent_pos"] = agent_data observation["agent_pos_mask"] = agent_pos_mask observation["dof_mask"] = dof_mask return observation def model_output_process(self, action_pred, state): if not self.args.delta_action and self.args.action_rpy: return action_pred pred_components = extract_components( action_pred, dim_dof_config, self.args.action_rpy ) pred_right_xyz = pred_components["right_xyz"] pred_right_rot = pred_components["right_rot"] pred_right_gripper = pred_components["right_gripper"] if "left_xyz" in pred_components: pred_left_xyz = pred_components["left_xyz"] pred_left_rot = pred_components["left_rot"] pred_left_gripper = pred_components["left_gripper"] else: pred_left_xyz = np.zeros((self.args.action_horizon, 3)) pred_left_rot = np.zeros((self.args.action_horizon, 3)) pred_left_gripper = np.zeros((self.args.action_horizon, 1)) post_action_pred = np.zeros((self.args.action_horizon, self.args.action_dim)) if self.args.delta_action: assert ( self.args.action_dim == 14 ), "Delta robot support and testing are not yet available." state_components = extract_components( state, dim_dof_config, self.args.state_rpy ) left_xyz = state_components["left_xyz"] left_rot = state_components["left_rot"] right_xyz = state_components["right_xyz"] right_rot = state_components["right_rot"] post_action_pred[:, :3] = pred_left_xyz + left_xyz post_action_pred[:, 3:6] = compose_state_and_delta_to_abs_rpy( pred_left_rot, left_rot[0] ) post_action_pred[:, 6:7] = pred_left_gripper post_action_pred[:, 7:10] = pred_right_xyz + right_xyz post_action_pred[:, 10:13] = compose_state_and_delta_to_abs_rpy( pred_right_rot, right_rot[0] ) post_action_pred[:, 13:14] = pred_right_gripper elif not self.args.action_rpy: post_action_pred[:, :3] = pred_left_xyz post_action_pred[:, 3:6] = so3_to_euler_zyx_batch_nb(pred_left_rot) post_action_pred[:, 6:7] = pred_left_gripper post_action_pred[:, 7:10] = pred_right_xyz post_action_pred[:, 10:13] = so3_to_euler_zyx_batch_nb(pred_right_rot) post_action_pred[:, 13:14] = pred_right_gripper else: post_action_pred = action_pred return post_action_pred def postprocess(self, action_pred, interpolate_multiplier=None): if interpolate_multiplier is None: interpolate_multiplier = self.args.interpolate_multiplier if isinstance(action_pred, torch.Tensor): action_pred = action_pred.to(torch.float32).cpu().squeeze(0).numpy() left_action_pred = action_pred[:, :7] # (32, 7) right_action_pred = action_pred[:, 7:14] # (32, 7) start_frame = int(self.args.action_start_ratio * len(left_action_pred)) end_frame = int(self.args.action_end_ratio * len(left_action_pred)) left_action_pred = left_action_pred[start_frame:end_frame] right_action_pred = right_action_pred[start_frame:end_frame] print("left_action_pred", left_action_pred[-1], flush=True) print("right_action_pred", right_action_pred[-1], flush=True) left_action_pred = left_action_pred.tolist() right_action_pred = right_action_pred.tolist() serialized_actions = { "follow1_pos": left_action_pred, "follow2_pos": right_action_pred, ## for joint-control # "follow1_joints":left_action_pred, # "follow2_joints":right_action_pred, } return serialized_actions def predict_action_rtc( self, state, views, instruction=None, valid_action_dim=7, update_subtask=False, action_predict_mode=None, ): if action_predict_mode is not None: self.action_predict_mode = action_predict_mode observation = self.preprocess(state, views, valid_action_dim) print("use instruction", instruction, flush=True) camera_names = [camera_to_view_mapping[key] for key in views.keys()] # camera_names = ["right_wrist_view", "global_view", "side_view"] print("mode:", self.action_predict_mode, flush=True) inputs = self._construct_input( observation, instruction, camera_names=camera_names, valid_action_dim=valid_action_dim, mode=self.action_predict_mode, ) model_action_dim = sum(self.dof_config.values()) padding = torch.zeros((1, model_action_dim)) norm_padding = self.normalizer_action.normalize_data( padding, [self.args.dataset_name] ) inputs["padding_action"] = norm_padding inputs = move_to_cuda(inputs, device="cuda:0") agent_data = observation["agent_pos"][..., : self.args.model_action_dim][0] print("before generate_flow_action", flush=True) print(inputs.keys(), flush=True) print(inputs["dataset_names"], flush=True) print(self.processor.tokenizer.decode(inputs["input_ids"][0]), flush=True) print(inputs["agent_pos_mask"][0], flush=True) print(inputs["dof_mask"][0, 0], flush=True) if self.action_predict_mode == "ar": action_pred = self.generate_ar_action(inputs) else: action_pred = self.generate_flow_action(inputs) print("after generate_flow_action", flush=True) if isinstance(action_pred, torch.Tensor): action_pred = action_pred.float().cpu().squeeze(0).numpy() if action_pred is None: return None if self.args.dct_scale > 0: scale = self.args.dct_scale dct_coeff = dct(action_pred, axis=0, norm="ortho") dct_coeff = np.around(dct_coeff * scale) action_pred = idct(dct_coeff / scale, axis=0, norm="ortho") if action_pred.shape[-1] == 7: print("Before concat action_pred", action_pred.shape, flush=True) right_action_pred = action_pred left_action_pred = np.zeros_like(right_action_pred) action_pred = np.concatenate([left_action_pred, right_action_pred], axis=1) print("After concat action_pred", action_pred.shape, flush=True) # unnorm action_pred # print("action_pred", action_pred[:, 3], flush=True) action_pred = ( self.normalizer_action.unnormalize_data( torch.tensor(action_pred).unsqueeze(0), [self.args.dataset_name] ) .squeeze(0) .numpy() ) print("After unnormalize_data action_pred", action_pred.shape, flush=True) action_pred = self.model_output_process(action_pred, agent_data) action_pred = self.postprocess(action_pred) return action_pred def generate_flow_action( self, inputs, last_action_chunk=None, max_guidance_weight=20.0, num_inference_timesteps=10, sigma_action=0.2, ): model_action_dim = sum(self.dof_config.values()) if last_action_chunk is None: output = self.model.generate_flow_action( action_horizon=self.args.action_horizon, action_dim=model_action_dim, num_inference_timesteps=num_inference_timesteps, unnorm=False, **inputs, ) else: output = self.model.generate_flow_action_rtc( action_horizon=self.args.action_horizon, action_dim=model_action_dim, num_inference_timesteps=num_inference_timesteps, inference_delay=self.args.rtc_inference_delay, execution_horizon=self.args.rtc_execution_horizon - 1, max_guidance_weight=max_guidance_weight, last_action_chunk=last_action_chunk, sigma_action=sigma_action, unnorm=False, **inputs, ) action_pred = output["predict_action"] # (b, action_horizon, action_dim) return action_pred def generate_text(self, inputs): return self.model.generate_text(**inputs) def _preprocess_ar_batch(self, batch): input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] moe_token_types = batch["moe_token_types"] labels = batch.get("labels", None) prefix_length = batch.get("prefix_length", None) generation_prompt_ids = torch.tensor( [151644, 77091], device=input_ids.device, dtype=input_ids.dtype ) # <|im_start|>assistant matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( input_ids[0, 1:] == generation_prompt_ids[1] ) if matches.any(): split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() # construct output ids gt_output_ids = input_ids[:, split_pos + 3 : prefix_length] # remove output part from input input_ids = input_ids[:, : split_pos + 3] moe_token_types = moe_token_types[:, : split_pos + 3] if attention_mask is not None: attention_mask = attention_mask[:, : split_pos + 3] if labels is not None: labels = labels[:, split_pos + 3 : prefix_length] batch.update( { "input_ids": input_ids, "attention_mask": attention_mask, "moe_token_types": moe_token_types, "labels": labels, "gt_output_ids": gt_output_ids, "prefix_length": split_pos + 3, } ) return batch def generate_ar_action(self, inputs): # batch = self._preprocess_ar_batch(batch=inputs) action_pred = None count = 0 while action_pred is None: if count > 5: # raise ValueError("re-generate ar action failed") return None count += 1 output = self.model.generate_ar_action( # action_dim=args.action_dim, action_dim=14, action_horizon=self.args.action_horizon, unnorm=False, **inputs, ) action_pred = output["predict_action"] action_pred = action_pred[0] return action_pred class WallxInfer: def __init__(self, args: WallxInferArgs): self.args = args self.model_wrapper = WallxModelWrapper(args) def run_infer_robochallenge( self, state, views, instruction, valid_action_dim=7, action_predict_mode=None ): action_pred = self.model_wrapper.predict_action_rtc( state=state, views=views, instruction=instruction, valid_action_dim=valid_action_dim, action_predict_mode=action_predict_mode, ) return action_pred if __name__ == "__main__": args = WallxInferArgs() Infer = WallxInfer(args) Infer.run_infer()