From a84756ff8743c49863f144b547abc2ea2bb4cd3d Mon Sep 17 00:00:00 2001 From: Trav1slaflame Date: Tue, 21 Oct 2025 22:13:20 +0800 Subject: [PATCH 1/6] adding vla support: Robovlms, SpatialVLA, and OpenVLA-OFT --- configs/policy/openvla_oft.yaml | 55 +++ configs/policy/robovlms.yaml | 33 ++ configs/policy/spatialvla.yaml | 33 ++ .../act_wrapper/robovlms_eepose_wrapper.py | 95 ++++ .../act_wrapper/spatialvla_eepose_wrapper.py | 95 ++++ .../obs_wrapper/robovlms_rgb_wrapper.py | 98 ++++ .../obs_wrapper/spatialvla_rgb_wrapper.py | 106 ++++ .../customize/policy_model/openvla_model.py | 329 +++++++++++++ .../customize/policy_model/robovlms_model.py | 452 ++++++++++++++++++ .../policy_model/spatialvla_model.py | 178 +++++++ maniunicon/policies/torch_model_vla.py | 336 +++++++++++++ maniunicon/utils/vla_utils.py | 28 ++ 12 files changed, 1838 insertions(+) create mode 100644 configs/policy/openvla_oft.yaml create mode 100644 configs/policy/robovlms.yaml create mode 100644 configs/policy/spatialvla.yaml create mode 100644 maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py create mode 100644 maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py create mode 100644 maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py create mode 100644 maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py create mode 100644 maniunicon/customize/policy_model/openvla_model.py create mode 100644 maniunicon/customize/policy_model/robovlms_model.py create mode 100644 maniunicon/customize/policy_model/spatialvla_model.py create mode 100644 maniunicon/policies/torch_model_vla.py create mode 100644 maniunicon/utils/vla_utils.py diff --git a/configs/policy/openvla_oft.yaml b/configs/policy/openvla_oft.yaml new file mode 100644 index 0000000..0ba2e58 --- /dev/null +++ b/configs/policy/openvla_oft.yaml @@ -0,0 +1,55 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 3 # expected openloop steps +use_real_time_chunking: false +enable_recording: false +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.openvla_model.OpenVLAModel + saved_model_path: null # path to openvla-oft ckpt + + device: ${device} + unnorm_key: "pick_up_anything_rlds/1.0.0" # choose the one that matching the task training dataset + + # Action configuration + use_act_chunk: true # Use action chunking (set to false for action ensemble) + action_ensemble_temp: -0.8 # Temperature for action ensemble (negative values give more weight to older predictions) + + # Model configuration + use_l1_regression: true # Use L1 regression action head for continuous actions + use_diffusion: false # Use diffusion action head (alternative to L1) + num_diffusion_steps_inference: 50 # Number of diffusion steps for inference (if use_diffusion=true) + + # Input configuration + use_proprio: false # Use proprioception input + center_crop: true # Apply center cropping to images + num_images_in_input: 1 # Number of camera views (1 for single camera, >1 for multi-camera) + + # Fine-tuning configuration + lora_rank: 32 # LoRA rank for fine-tuning + + # Quantization options (for memory efficiency) + load_in_8bit: false # Load model with 8-bit quantization + load_in_4bit: false # Load model with 4-bit quantization + + # Task specification + task_name: "pick up anything" # Task instruction for the VLA model + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.spatialvla_rgb_wrapper.SpatialVLAImageWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + num_cams: 1 # Set to 1 for single camera, 2+ for multi-camera setup + observation_horizon: 1 # Should match the window size in VLA model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.spatialvla_eepose_wrapper.SpatialVLAEEPoseWrapper + hist_action: 0 # Number of historical actions to consider + action_horizon: 5 # Should match NUM_ACTIONS_CHUNK for the robot platform (8 for LIBERO, 25 for ALOHA, 5 for Bridge, 5 for xarm sft) + control_mode: cartesian + dt: ${policy.dt} diff --git a/configs/policy/robovlms.yaml b/configs/policy/robovlms.yaml new file mode 100644 index 0000000..8b224d3 --- /dev/null +++ b/configs/policy/robovlms.yaml @@ -0,0 +1,33 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 3 # expected openloop steps +use_real_time_chunking: false +enable_recording: false +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.robovlms_model.RoboVlmsModel + ckpt_path: null # path to robovlms model ckpt + config_path: null # path to robovlms model config file + device: ${device} + log_save_dir: null # path to save ckpt loading msg, defult is null + use_act_chunk: true + task_name: "pick up anything" # task instruction that input to vla model + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.robovlms_rgb_wrapper.RoboVlmsImageWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + num_cams: 2 + observation_horizon: 1 # should be the same as the window size in vla model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.robovlms_eepose_wrapper.RoboVlmsEEPoseWrapper + hist_action: 0 # number of historical actions to consider + action_horizon: 5 # should be the same as the action chunk size in vla model + control_mode: cartesian + dt: ${policy.dt} \ No newline at end of file diff --git a/configs/policy/spatialvla.yaml b/configs/policy/spatialvla.yaml new file mode 100644 index 0000000..21fcec4 --- /dev/null +++ b/configs/policy/spatialvla.yaml @@ -0,0 +1,33 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 3 # expected openloop steps +use_real_time_chunking: false +enable_recording: false +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.spatialvla_model.SpatialVLAModel + saved_model_path: null # path to spatialvla ckpt + device: ${device} + unnorm_key: "pick_up_anything_rlds/1.0.0" # choose the one that matching the task training dataset + use_act_chunk: true # set to true to use action chunking, else action ensemble + action_ensemble_temp: -0.8 + task_name: "pick up anything" # task instruction that input to vla model + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.spatialvla_rgb_wrapper.SpatialVLAImageWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + num_cams: 2 + observation_horizon: 1 # should be the same as the window size in vla model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.spatialvla_eepose_wrapper.SpatialVLAEEPoseWrapper + hist_action: 0 # number of historical actions to consider + action_horizon: 5 # should be the same as the action chunk size in vla model + control_mode: cartesian + dt: ${policy.dt} \ No newline at end of file diff --git a/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py b/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py new file mode 100644 index 0000000..62fd5f1 --- /dev/null +++ b/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py @@ -0,0 +1,95 @@ +import time +import torch +import numpy as np +from typing import Union + +from maniunicon.utils.shared_memory.shared_storage import RobotAction + + +class RoboVlmsEEPoseWrapper: + """ + A wrapper class for handling action for Robovlms. + robot control mode: cartesian/ee pose + """ + + def __init__( + self, + action_horizon: int = 8, + hist_action: int = 0, + action_exec_latency: float = 0.01, + control_mode: str = "cartesian", + synchronized: bool = False, + dt: float = 0.02, + **kwargs, + ): + """ + Initialize the wrapper with history action and open loop steps. + :param action_horizon: Number of open loop steps to execute. + :param hist_action: Number of historical actions to consider. + """ + self.action_horizon = action_horizon + self.hist_action = hist_action + self.action_exec_latency = action_exec_latency + self.control_mode = control_mode + self.synchronized = synchronized + self.dt = dt + + def __call__( + self, + actions: Union[torch.Tensor, np.ndarray], + timestamp: float, + start_timestamp: float, + return_raw_actions: bool = True, + ): + robot_actions = [] + action_timestamps = np.arange(self.action_horizon) * self.dt + timestamp + curr_time = time.time() + # if synchronized execution is enabled, we do not need to check the action timestamp + if self.synchronized: + is_new = np.ones(action_timestamps.shape[0], dtype=bool) + else: + is_new = action_timestamps > (curr_time + self.action_exec_latency) + + if isinstance(actions, torch.Tensor): + actions = actions.cpu().numpy() + + raw_actions = actions.copy() + if actions.ndim == 4: + # action: (bs, ws, fwd, action_dim) remove the batch dimension and history actions + actions = actions[0, 0, self.hist_action : self.hist_action + self.action_horizon, :] + elif actions.ndim == 2: + # action: (fwd, action_dim) + actions = actions[self.hist_action : self.hist_action + self.action_horizon, :] + else: + raise ValueError(f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array.") + + print(np.sum(is_new), " new actions, ", len(actions), " total actions") + if np.sum(is_new) == 0: + # exceeded time budget, still do something + actions = actions[-1:, :] + next_step_idx = int(np.ceil((curr_time - start_timestamp) / self.dt)) + action_timestamp = start_timestamp + next_step_idx * self.dt + print("No new actions !!! Check model inference time") + action_timestamps = np.array([action_timestamp]) + else: + actions = actions[is_new, :] + action_timestamps = action_timestamps[is_new] + + for idx in range(len(actions)): + action = actions[idx, :] + # Convert to RobotAction + robot_actions.append( + RobotAction( + # FK is computed by the robot + tcp_position=action[:3], + tcp_orientation=action[3:7], + gripper_state=np.array([(action[7:8] > 0.5).astype(int)]), + control_mode=self.control_mode, + timestamp=timestamp, + target_timestamp=action_timestamps[idx], + ) + ) + if return_raw_actions: + return robot_actions, raw_actions + else: + return robot_actions diff --git a/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py b/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py new file mode 100644 index 0000000..aac8ef3 --- /dev/null +++ b/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py @@ -0,0 +1,95 @@ +import time +import torch +import numpy as np +from typing import Union + +from maniunicon.utils.shared_memory.shared_storage import RobotAction + + +class SpatialVLAEEPoseWrapper: + """ + A wrapper class for handling action for Robovlms. + robot control mode: cartesian/ee pose + """ + + def __init__( + self, + action_horizon: int = 8, + hist_action: int = 0, + action_exec_latency: float = 0.01, + control_mode: str = "cartesian", + synchronized: bool = False, + dt: float = 0.02, + **kwargs, + ): + """ + Initialize the wrapper with history action and open loop steps. + :param action_horizon: Number of open loop steps to execute. + :param hist_action: Number of historical actions to consider. + """ + self.action_horizon = action_horizon + self.hist_action = hist_action + self.action_exec_latency = action_exec_latency + self.control_mode = control_mode + self.synchronized = synchronized + self.dt = dt + + def __call__( + self, + actions: Union[torch.Tensor, np.ndarray], + timestamp: float, + start_timestamp: float, + return_raw_actions: bool = True, + ): + robot_actions = [] + action_timestamps = np.arange(self.action_horizon) * self.dt + timestamp + curr_time = time.time() + # if synchronized execution is enabled, we do not need to check the action timestamp + if self.synchronized: + is_new = np.ones(action_timestamps.shape[0], dtype=bool) + else: + is_new = action_timestamps > (curr_time + self.action_exec_latency) + + if isinstance(actions, torch.Tensor): + actions = actions.cpu().numpy() + + raw_actions = actions.copy() + if actions.ndim == 4: + # action: (bs, ws, fwd, action_dim) remove the batch dimension and history actions + actions = actions[0, 0, self.hist_action : self.hist_action + self.action_horizon, :] + elif actions.ndim == 2: + # action: (fwd, action_dim) + actions = actions[self.hist_action : self.hist_action + self.action_horizon, :] + else: + raise ValueError(f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array.") + + print(np.sum(is_new), " new actions, ", len(actions), " total actions") + if np.sum(is_new) == 0: + # exceeded time budget, still do something + actions = actions[-1:, :] + next_step_idx = int(np.ceil((curr_time - start_timestamp) / self.dt)) + action_timestamp = start_timestamp + next_step_idx * self.dt + print("No new actions !!! Check model inference time") + action_timestamps = np.array([action_timestamp]) + else: + actions = actions[is_new, :] + action_timestamps = action_timestamps[is_new] + + for idx in range(len(actions)): + action = actions[idx, :] + # Convert to RobotAction + robot_actions.append( + RobotAction( + # FK is computed by the robot + tcp_position=action[:3], + tcp_orientation=action[3:7], + gripper_state=np.array([(action[7:8] > 0.5).astype(int)]), + control_mode=self.control_mode, + timestamp=timestamp, + target_timestamp=action_timestamps[idx], + ) + ) + if return_raw_actions: + return robot_actions, raw_actions + else: + return robot_actions diff --git a/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py new file mode 100644 index 0000000..87d25f7 --- /dev/null +++ b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py @@ -0,0 +1,98 @@ +from typing import List + +import cv2 +import torch +import numpy as np + + +def resize_image_sequence(images, target_size, interp=cv2.INTER_AREA): + """ + Resize an image sequence using OpenCV + + Args: + images: numpy array of shape (N, H, W, C) where: + N = number of images + H = height + W = width + C = channels + target_size: tuple of (height, width) + + Returns: + resized images array of shape (N, new_H, new_W, C) + """ + N, H, W, C = images.shape + new_H, new_W = target_size + + # Reshape to 2D array of images for faster processing + reshaped = images.reshape(-1, H, W, C) + + # Preallocate output array + output = np.empty((N, new_H, new_W, C), dtype=images.dtype) + + # Resize each image + for i in range(N): + res = cv2.resize( + images[i], (new_W, new_H), interpolation=interp + ) + if C == 1: + output[i] = res[:, :, np.newaxis] + else: + output[i] = res + + return output + + +class RoboVlmsImageWrapper: + def __init__( + self, + camera_config, + state_keys: List[str], + num_cams: int = 1, + shared_storage=None, + device: str = "cpu", + observation_horizon: int = 1, + ): + self.camera_config = camera_config + self.state_keys = state_keys + self.num_cams = num_cams + self.shared_storage = shared_storage + self.device = device + + def __call__(self, state, camera): + """ + Wraps the observation for the RoboVlms model. + :return: Wrapped observation tensor. + """ + assert state is not None, "State cannot be None" + assert camera is not None, "Camera cannot be None" + + state_tensor = ( + torch.from_numpy( + np.concatenate( + [getattr(state, key) for key in self.state_keys], axis=-1 + ) + ) + .to(self.device) + .float() + ) + + colors = resize_image_sequence( + camera.colors.reshape(-1, *camera.colors.shape[2:]), (224, 224) + ) + colors = colors.reshape( + camera.colors.shape[0], + camera.colors.shape[1], + *colors.shape[1:], + ) + + images_tensor = {} + for cam_idx in range(colors.shape[1]): + images_tensor[f"camera_{cam_idx}"] = torch.from_numpy( + colors[:, cam_idx] + ).to(self.device) + + obs_tensor = { + "state": state_tensor, + "image": images_tensor, + } + return obs_tensor diff --git a/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py b/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py new file mode 100644 index 0000000..73625dd --- /dev/null +++ b/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py @@ -0,0 +1,106 @@ +from typing import List + +from PIL import Image +import torch +import numpy as np + + +def resize_image_sequence(images, target_size, resample=Image.BICUBIC): + """ + Resize an image sequence using PIL to match training data processing when converting from zarr to rlds + + Args: + images: numpy array of shape (N, H, W, C) where: + N = number of images + H = height + W = width + C = channels + target_size: int or tuple of (height, width) + if int, both height and width will be set to this value + resample: PIL resampling filter (default: Image.BICUBIC) + + Returns: + resized images array of shape (N, new_size, new_size, 3) + Note: output will always have 3 channels (RGB) + """ + # Convert target_size to tuple if it's a single integer + if isinstance(target_size, int): + target_size = (target_size, target_size) + new_H, new_W = target_size + + N, H, W, C = images.shape + + # Preallocate output array (always 3 channels for RGB) + output = np.empty((N, new_H, new_W, 3), dtype=np.uint8) + + # Resize each image using PIL + for i in range(N): + # Convert numpy array to PIL Image + img_pil = Image.fromarray(images[i]) + + # Convert to RGB (regardless of original number of channels) + img_pil = img_pil.convert("RGB") + + # Resize image + img_resized = img_pil.resize((new_W, new_H), resample=resample) + + # Convert back to numpy array and store in output + output[i] = np.array(img_resized) + + return output + + +class SpatialVLAImageWrapper: + def __init__( + self, + camera_config, + state_keys: List[str], + num_cams: int = 1, + shared_storage=None, + device: str = "cpu", + observation_horizon: int = 1, + ): + self.camera_config = camera_config + self.state_keys = state_keys + self.num_cams = num_cams + self.shared_storage = shared_storage + self.device = device + + def __call__(self, state, camera): + """ + Wraps the observation for the RoboVlms model. + :return: Wrapped observation tensor. + """ + assert state is not None, "State cannot be None" + assert camera is not None, "Camera cannot be None" + + state_tensor = ( + torch.from_numpy( + np.concatenate( + [getattr(state, key) for key in self.state_keys], axis=-1 + ) + ) + .to(self.device) + .float() + ) + + colors = resize_image_sequence( + camera.colors.reshape(-1, *camera.colors.shape[2:]), (224, 224) + ) + colors = colors.reshape( + camera.colors.shape[0], + camera.colors.shape[1], + *colors.shape[1:], + ) + + images_tensor = {} + for cam_idx in range(colors.shape[1]): + images_tensor[f"camera_{cam_idx}"] = torch.from_numpy( + colors[:, cam_idx] + ).to(self.device) + + obs_tensor = { + "state": state_tensor, + "image": images_tensor, + } + return obs_tensor diff --git a/maniunicon/customize/policy_model/openvla_model.py b/maniunicon/customize/policy_model/openvla_model.py new file mode 100644 index 0000000..97041e8 --- /dev/null +++ b/maniunicon/customize/policy_model/openvla_model.py @@ -0,0 +1,329 @@ +from typing import Optional, List, Dict, Any, Union +import os +import numpy as np +from collections import deque +from PIL import Image +import torch +import cv2 as cv +import time +from pathlib import Path +from dataclasses import dataclass + +# Import OpenVLA components from the installed package +from experiments.robot.openvla_utils import ( + get_vla, + get_processor, + get_action_head, + get_proprio_projector, + get_vla_action, + prepare_images_for_vla, + DEVICE, +) +from experiments.robot.robot_utils import get_image_resize_size +from prismatic.vla.constants import ACTION_DIM, PROPRIO_DIM + + +class ActionEnsembler: + def __init__(self, pred_action_horizon, action_ensemble_temp=0.0): + self.pred_action_horizon = pred_action_horizon + self.action_ensemble_temp = action_ensemble_temp + self.action_history = deque(maxlen=self.pred_action_horizon) + + def reset(self): + self.action_history.clear() + + def ensemble_action(self, cur_action): + self.action_history.append(cur_action) + num_actions = len(self.action_history) + if cur_action.ndim == 1: + curr_act_preds = np.stack(self.action_history) + else: + curr_act_preds = np.stack( + [ + pred_actions[i] + for (i, pred_actions) in zip( + range(num_actions - 1, -1, -1), self.action_history + ) + ] + ) + # if temp > 0, more recent predictions get exponentially *less* weight than older predictions + weights = np.exp(-self.action_ensemble_temp * np.arange(num_actions)) + weights = weights / weights.sum() + # compute the weighted average across all predictions for this timestep + cur_action = np.sum(weights[:, None] * curr_act_preds, axis=0) + + return cur_action + + +@dataclass +class OpenVLAConfig: + """Configuration for OpenVLA model wrapper""" + + # Model configuration + pretrained_checkpoint: Union[str, Path] = "openvla/openvla-7b" + model_family: str = "openvla" + + # Action head configuration + use_l1_regression: bool = False + use_diffusion: bool = False + num_diffusion_steps_train: int = 50 + num_diffusion_steps_inference: int = 50 + + # Input configuration + use_film: bool = False + num_images_in_input: int = 1 + use_proprio: bool = False + center_crop: bool = False + + # LoRA configuration + lora_rank: int = 32 + + # Normalization + unnorm_key: str = "" + use_relative_actions: bool = False + + # Quantization + load_in_8bit: bool = False + load_in_4bit: bool = False + + # Other + seed: int = 7 + + +class OpenVLAModel: + def __init__( + self, + device, + saved_model_path: str = "openvla/openvla-7b", + unnorm_key: str = None, + image_size: List[int] = [224, 224], + action_ensemble_temp: float = 0.0, + use_act_chunk: bool = True, + task_name: str = None, + use_l1_regression: bool = False, + use_diffusion: bool = False, + num_diffusion_steps_inference: int = 50, + use_proprio: bool = False, + center_crop: bool = False, + num_images_in_input: int = 1, + lora_rank: int = 32, + load_in_8bit: bool = False, + load_in_4bit: bool = False, + ) -> None: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + self.device = device + self.unnorm_key = unnorm_key + print(f"*** unnorm_key: {unnorm_key} ***") + + # Setup task instruction + self.task_name = task_name + + # Create configuration object + self.cfg = OpenVLAConfig( + pretrained_checkpoint=saved_model_path, + unnorm_key=unnorm_key, + use_l1_regression=use_l1_regression, + use_diffusion=use_diffusion, + num_diffusion_steps_inference=num_diffusion_steps_inference, + use_proprio=use_proprio, + center_crop=center_crop, + num_images_in_input=num_images_in_input, + lora_rank=lora_rank, + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + ) + + # Load VLA model + print("Loading OpenVLA model...") + self.vla = get_vla(self.cfg) + + # Load processor + self.processor = get_processor(self.cfg) + + # Load proprioception projector if needed + self.proprio_projector = None + if self.cfg.use_proprio: + self.proprio_projector = get_proprio_projector( + self.cfg, self.vla.llm_dim, PROPRIO_DIM + ) + + # Load action head if needed + self.action_head = None + if self.cfg.use_l1_regression or self.cfg.use_diffusion: + self.action_head = get_action_head(self.cfg, self.vla.llm_dim) + + # Image configuration + self.image_size = image_size + self.resize_size = get_image_resize_size(self.cfg) + + # Get observation horizon from processor if available + if hasattr(self.processor, "num_obs_steps") and hasattr( + self.processor, "obs_delta" + ): + self.obs_horizon = ( + self.processor.num_obs_steps - 1 + ) * self.processor.obs_delta + 1 + self.obs_interval = self.processor.obs_delta + else: + # Default values if not available + self.obs_horizon = 1 + self.obs_interval = 1 + + # Action chunking configuration + if hasattr(self.processor, "action_chunk_size"): + self.pred_action_horizon = self.processor.action_chunk_size + else: + # Default based on constants + from prismatic.vla.constants import NUM_ACTIONS_CHUNK + + self.pred_action_horizon = NUM_ACTIONS_CHUNK + + self.image_history = deque(maxlen=self.obs_horizon) + self.use_act_chunk = use_act_chunk + + # Action ensemble configuration + if self.use_act_chunk: + action_ensemble = False + else: + action_ensemble = True + self.action_ensemble = action_ensemble + self.action_ensemble_temp = action_ensemble_temp + + if self.action_ensemble: + self.action_ensembler = ActionEnsembler( + self.pred_action_horizon, self.action_ensemble_temp + ) + else: + self.action_ensembler = None + + self.rollout_step_counter = 0 + + def reset(self) -> None: + print("reset model now!!!!!!!!") + self.image_history.clear() + if self.action_ensemble: + self.action_ensembler.reset() + self.rollout_step_counter = 0 + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + image = cv.resize(image, tuple(self.image_size), interpolation=cv.INTER_AREA) + return image + + def _add_image_to_history(self, image: np.ndarray) -> None: + if len(self.image_history) == 0: + self.image_history.extend([image] * self.obs_horizon) + else: + self.image_history.append(image) + + def _obtain_image_history(self) -> List[Image.Image]: + image_history = list(self.image_history) + images = image_history[:: self.obs_interval] + images = [Image.fromarray(image).convert("RGB") for image in images] + return images + + def preprocess(self, obs): + """Preprocess observation from ManiUniCon format to OpenVLA format""" + # Extract primary camera image + obs["image"]["camera_1"] = ( + obs["image"]["camera_1"].squeeze().cpu().detach() + ) # (224, 224, 3) + + # Process primary image + image = obs["image"]["camera_1"].numpy() + assert image.dtype == np.uint8 + image = self._resize_image(image) + + # For OpenVLA, we'll return the image directly for processing + # The history handling depends on the specific OpenVLA configuration + if self.obs_horizon > 1: + self._add_image_to_history(image) + images = self._obtain_image_history() + else: + images = [Image.fromarray(image).convert("RGB")] + + # Prepare observation dict for OpenVLA + obs_dict = { + "full_image": image, # Primary camera as numpy array + } + + # Add wrist cameras if configured for multi-image input + if self.cfg.num_images_in_input > 1: + # Check if there are additional camera views in the observation + for key in obs["image"].keys(): + if "wrist" in key.lower() or "camera_0" in key: + wrist_img = obs["image"][key].squeeze().cpu().detach().numpy() + wrist_img = self._resize_image(wrist_img) + obs_dict[key] = wrist_img + + # Add proprioception state if available and configured + if self.cfg.use_proprio and "state" in obs: + if isinstance(obs["state"], torch.Tensor): + obs_dict["state"] = obs["state"].cpu().numpy() + else: + obs_dict["state"] = obs["state"] + + # Add task instruction + obs_dict["instruction"] = ( + self.task_name if self.task_name else "complete the task" + ) + + return obs_dict + + def __call__(self, obs, **kwargs): + """ + Forward pass through the model. + :param obs: Observation input to the model. + :return: Model output. + """ + # Preprocess observation to OpenVLA format + obs_dict = self.preprocess(obs) + + # Get action from OpenVLA + start = time.time() + action_list = get_vla_action( + self.cfg, + self.vla, + self.processor, + obs_dict, + obs_dict["instruction"], + action_head=self.action_head, + proprio_projector=self.proprio_projector, + use_film=self.cfg.use_film, + ) + print(f"**** OpenVLA-OFT inference time: {time.time() - start}") + + # Convert list of actions to numpy array + raw_actions = np.array(action_list) + + # Handle action ensemble if configured + if self.action_ensemble: + print("ensemble action!!!") + raw_actions = self.action_ensembler.ensemble_action(raw_actions) + if raw_actions.ndim == 1: + raw_actions = raw_actions[None] + + self.rollout_step_counter += 1 + + # Ensure we have numpy array + if isinstance(raw_actions, torch.Tensor): + raw_actions = raw_actions.cpu().numpy() + + # Convert rotation from euler to quaternion + # OpenVLA outputs: [x, y, z, rx, ry, rz, gripper] + # Need to convert to: [x, y, z, qx, qy, qz, qw, gripper] + from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = raw_actions.shape + if A == 7: # Euler format + raw_actions = np.concatenate( + [ + euler_pose_to_quat( + raw_actions[..., :-1].reshape(-1, A - 1) + ).reshape(N, -1), + raw_actions[..., -1:], + ], + axis=-1, + ) + + print(f"step {self.rollout_step_counter} action {raw_actions}") + + return raw_actions diff --git a/maniunicon/customize/policy_model/robovlms_model.py b/maniunicon/customize/policy_model/robovlms_model.py new file mode 100644 index 0000000..9cde5e7 --- /dev/null +++ b/maniunicon/customize/policy_model/robovlms_model.py @@ -0,0 +1,452 @@ +import json +import os.path +from copy import deepcopy +import torch +from PIL import Image +from typing import Literal +import numpy as np +import functools +import time + +# Import RoboVLMs components from the installed package +from robovlms.train.base_trainer import BaseTrainer +from robovlms.utils.model_utils import build_tokenizer +from robovlms.data.data_utils import get_text_function +from robovlms.data.data_utils import ( + preprocess_image, + get_prompt_builder, + tcp_to_world_frame, +) +from robovlms.utils.config_utils import load_config +from queue import Queue +from robovlms.model.policy_head.action_tokenizer import ActionTokenizer + +fwd_decay_ratio = 1 + +class RoboVlmsModel: + # model option + def __init__( + self, + ckpt_path, + config_path, + device, + log_save_dir=None, + raw_calvin=True, + debug=False, + use_act_chunk=False, + task_name=None + ): + # setup the task instruction that input to vla model + self.task_name = task_name + # Loading robovlms model configs + assert config_path != None + configs = load_config(config_path) + self.model = BaseTrainer(configs=configs) + + # Get checkpoint path + print("ckpt_path", ckpt_path) + from robovlms.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict + # Handle DeepSpeed ckpt + if os.path.isdir(ckpt_path): + target_ckpt_path = ckpt_path.replace(".ckpt", ".pt") + print(f"converting {ckpt_path} to {target_ckpt_path}") + convert_zero_checkpoint_to_fp32_state_dict(ckpt_path, target_ckpt_path) + ckpt_path = target_ckpt_path + + self.init_config(ckpt_path, configs, device, log_save_dir, raw_calvin, debug, use_act_chunk) + # self.model.model.lm_head.window_size = 1 + + def init_config( + self, ckpt_path, configs, device, save_dir=None, raw_calvin=False, debug=False, use_act_chunk=False + ): + ### load and convert checkpoint + self.debug = debug + self.use_act_chunk = use_act_chunk + if configs["model"] == "kosmos": + import transformers + + package_dir = transformers.__path__[0] + os.system( + "cp /tools/modeling_kosmos2.py {}/models/kosmos2/modeling_kosmos2.py".format( + package_dir + ) + ) + + if not self.debug: + ckpt = torch.load(ckpt_path, map_location="cpu") + if "state_dict" in ckpt: + new_state_dict = ckpt["state_dict"] + elif "model_state_dict" in ckpt: + new_state_dict = ckpt["model_state_dict"] + else: + raise KeyError("no checkpoint dict in the loaded pretrain parameters") + + new_state_dict = self.convert_old_state_dict(new_state_dict) + msg = self.model.load_state_dict(new_state_dict, strict=False) + print(f"RoboVLMs CKPT Loaded \n {msg}") + del new_state_dict + + ckpt_dir = os.path.dirname(ckpt_path) + ckpt_name = os.path.basename(ckpt_path) + save_dir = ckpt_dir if save_dir is None else save_dir + load_info_path = os.path.join(save_dir, f"{ckpt_name}_loading_msg.json") + if os.path.exists(load_info_path): + os.system(f"rm {load_info_path}") + with open(load_info_path, "w") as f: + _info = { + "missing_keys": msg.missing_keys, + "unexpected_keys": msg.unexpected_keys, + } + json.dump(_info, f, indent=2) + print(f"Model loading msg is updated to: {load_info_path}") + + self.configs = configs + + dtype = torch.float32 + if self.configs["trainer"]["precision"] == "bf16": + dtype = torch.bfloat16 + elif self.configs["trainer"]["precision"] == "fp16": + dtype = torch.float16 + self.dtype = dtype + self.act_head_configs = self.configs["act_head"] + self.raw_calvin = raw_calvin + self.tcp_rel = self.configs.get("tcp_rel", False) + + print(f"raw action: {self.raw_calvin}") + + self.device = device + self.policy = self.model + self.policy = self.policy.to(self.dtype) + # self.policy = self.policy.float() + self.policy.to(self.device) + self.policy.eval() + + if not hasattr(self.policy.model, "lm_head"): + self.policy.model.lm_head = self.policy.model.act_head + + self.tokenizer = build_tokenizer(self.configs["tokenizer"]) + + self.window_size = configs["window_size"] + self.fwd_pred_next_n = configs["fwd_pred_next_n"] + self.act_step = self.fwd_pred_next_n + 1 + self.seq_len = self.configs["seq_len"] + self.use_hand_rgb = self.configs["use_hand_rgb"] + + if hasattr(self, "policy_setup"): + data_mix = "bridge" if self.policy_setup == "widowx_bridge" else "rt_1" + configs["train_dataset"]["data_mix"] = data_mix + configs["val_dataset"]["data_mix"] = data_mix + + image_preprocess = self.model.model.image_processor + self.image_preprocess = functools.partial( + preprocess_image, + image_processor=image_preprocess, + model_type=configs["model"], + ) + + self.text_preprocess = get_text_function( + self.model.model.tokenizer, configs["model"] + ) + + self.action_space = self.configs["act_head"].get("action_space", "continuous") + if self.action_space == "discrete": + self.action_tokenizer = ActionTokenizer( + self.tokenizer, + bins=self.act_head_configs["n_bin"], + min_action=self.act_head_configs["min_action"], + max_action=self.act_head_configs["max_action"], + ) + + print(f"Evaluating checkpoint {ckpt_path}") + + self.rgb_list = [] + self.hand_rgb_list = [] + self.action_hist_list = [] + self.rollout_step_counter = 0 + + self.vision_queue = Queue(maxsize=self.window_size) + self.vision_gripper_queue = Queue(maxsize=self.window_size) + self.action_queue = Queue(maxsize=self.window_size - 1) + + def ensemble_action(self, action): + if action.ndim >= 3: + action = action.squeeze() + + if action.ndim == 1: + action = action.unsqueeze(0) + + self.action_hist_list.append(action) + + act_cache = [] + max_len = 5 + while len(self.action_hist_list) > max_len: + self.action_hist_list.pop(0) + + idx = 0 + for act in self.action_hist_list[::-1]: + act_cache.append(act[idx]) + idx += 1 + + act_cache = torch.stack(act_cache, dim=0) + + weights = torch.tensor([fwd_decay_ratio**i for i in range(len(act_cache))]) + weights = weights / weights.sum() + + weighted_act = (act_cache * weights.unsqueeze(1)).sum(dim=0) + + return weighted_act + + @staticmethod + def convert_old_state_dict(state_dict): + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("module."): + new_k = k.replace("module.", "") + else: + new_k = k + + if not new_k.startswith("model."): + new_k = "model." + new_k + + new_state_dict[new_k] = state_dict[k] + return new_state_dict + + def _get_default_calvin_config(self): + return { + "type": "DiskCalvinDataset", + "data_dir": "CALVIN/task_ABCD_D/val", + "c_act_scaler": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + } + + def add_element_to_queue(self, q: Queue, element): + while q.qsize() >= q.maxsize: + q.get() + q.put(element) + + def get_history(self, q: Queue, pad: Literal["zero", "first"] = "zero"): + queue_list = list(q.queue) + if len(queue_list) == 0: + return queue_list, None + history_type = self.configs["act_head"].get("history_type", "pre") + if history_type == "pre": + pad_len = 0 + else: + raise ValueError(f"Unsupported history type {history_type}") + element = queue_list[0] + if pad == "zero": + if isinstance(element, torch.Tensor): + element = torch.zeros_like(element) + elif isinstance(element, np.ndarray): + element = np.zeros_like(element) + else: + raise ValueError("This type is not supported") + queue_list = [element for _ in range(pad_len)] + queue_list + else: + if isinstance(element, torch.Tensor): + pad_list = [element.clone() for _ in range(pad_len)] + elif isinstance(element, np.ndarray): + pad_list = [deepcopy(element) for _ in range(pad_len)] + queue_list = pad_list + queue_list + pad_mask = np.ones(q.maxsize, dtype=bool) + pad_mask[:pad_len] = False + return queue_list, pad_mask + + def preprocess(self, obs, lang, mode="continuous"): + obs["image"]["camera_0"] = obs["image"]["camera_0"].squeeze().cpu().detach() # (224, 224, 3) + obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() + # preprocess static cam image + image = obs["image"]["camera_1"].numpy() + image = Image.fromarray(image) + image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) + + gripper_x = None + if "camera_0" in obs["image"]: + gripper = obs["image"]["camera_0"].numpy() + gripper = Image.fromarray(gripper) + gripper_x = self.image_preprocess([gripper]).unsqueeze(0) + gripper_x = gripper_x.to(self.device).to(self.dtype) + + if self.configs["act_head"].get("history_type", "post") == "pre": + self.add_element_to_queue(self.vision_queue, image_x) + image_x, _ = self.get_history(self.vision_queue, pad="first") + image_x = torch.concatenate(image_x, dim=1) + + if gripper_x is not None: + self.add_element_to_queue(self.vision_gripper_queue, gripper_x) + gripper_x, _ = self.get_history(self.vision_gripper_queue, pad="first") + gripper_x = ( + torch.concatenate(gripper_x, dim=1).to(self.device).to(self.dtype) + ) + + if mode == "discrete": + if "llava" in self.policy.configs: + model_name = self.policy.configs["llava"] + elif "qwen" in self.policy.configs: + model_name = "qwen" + else: + # model_name = self.policy.configs['llm']['pretrained_model_name_or_path'] + model_name = self.policy.configs["model"] + + prompt_builder = get_prompt_builder( + model_name, bos=self.tokenizer.bos_token, eos=self.tokenizer.eos_token + ) + + conversation = [ + { + "from": "human", + "value": ( + f"What action should the robot take to {lang}?" + if self.act_step == 1 + else f"What {self.act_step} step actions should the robot take to {lang}?" + ), + }, + {"from": "gpt", "value": ""}, + ] + + input_ids = [] + for turn in conversation: + prompt_builder.add_turn(turn["from"], turn["value"]) + + input_ids = torch.tensor( + list( + self.tokenizer( + prompt_builder.get_prompt(), add_special_tokens=True + ).input_ids + ) + ) + if self.tokenizer.eos_token is not None: + input_ids = input_ids[:-1] + + text_x = input_ids.unsqueeze(0) + mask = torch.full((1, text_x.shape[-1]), True, dtype=torch.bool) + else: + text_x, mask = self.text_preprocess([lang]) + + return ( + image_x.to(self.device).to(self.dtype), + gripper_x, + text_x.to(self.device), + mask.to(self.device), + ) + + def reset(self): + print("reset model now!!!!!!!!") + if hasattr(self.model.model, "lm_head"): + self.model.model.lm_head.hidden_state = None + self.model.model.lm_head.history_memory = [] + if hasattr(self.model.model, "act_head"): + self.model.model.act_head.hidden_state = None + self.model.model.act_head.history_memory = [] + + self.rgb_list = [] + self.hand_rgb_list = [] + self.rollout_step_counter = 0 + self.action_hist_list = [] + + while not self.vision_queue.empty(): + self.vision_queue.get() + while not self.vision_gripper_queue.empty(): + self.vision_gripper_queue.get() + while not self.action_queue.empty(): + self.action_queue.get() + + def __call__(self, obs, **kwargs): + """ + Forward pass through the model. + :param obs: Observation input to the model. + :return: Model output. + """ + """Step function.""" + input_dict = dict() + image_x, gripper_x, text_x, mask = self.preprocess(obs, self.task_name, self.action_space) + + input_dict["rgb"] = image_x + input_dict["hand_rgb"] = gripper_x + input_dict["text"] = text_x + input_dict["text_mask"] = mask + + if self.action_space == "discrete": + input_dict["instr_and_action_ids"] = text_x + input_dict["instr_and_action_mask"] = mask + + start = time.time() + with torch.no_grad(): + action = self.policy.inference_step(input_dict)["action"] + # print("original action from model: ", action) + print("**** RoboVLMs inference time: ", time.time() - start) + + if self.action_space != "discrete": + if action[0].ndim == action[1].ndim + 1: + action = (action[0], action[1].unsqueeze(2)) + action = torch.cat( + [action[0], (torch.nn.functional.sigmoid(action[1]) > 0.5).float()], + dim=-1, + ) + + if isinstance(action, tuple): + action = torch.cat([action[0], action[1]], dim=-1) + + if isinstance(action, np.ndarray): + action = torch.from_numpy(action) + + if action.ndim == 2: + action = action.unsqueeze(1) + + if action.ndim == 3: + action = action.unsqueeze(1) + + action = action.detach().cpu() + + if self.tcp_rel: + robot_obs = ( + torch.from_numpy(obs["robot_obs"]) + .unsqueeze(0) + .unsqueeze(0) + .unsqueeze(0) + .repeat(1, self.window_size, self.fwd_pred_next_n, 1) + ) + action = tcp_to_world_frame(action, robot_obs) + + if not self.use_act_chunk: + print("ensemble action!!!") + action = self.ensemble_action(action) + + if isinstance(action, torch.Tensor): + action = action.squeeze() # (fwd, 7) for action chunk, (7,) for ensembled action + if not self.use_act_chunk: + action = action.unsqueeze(0) + + if self.configs.get("use_mu_law", False): + from robovlms.data.data_utils import inverse_mu_law_companding + + action = inverse_mu_law_companding( + action, self.configs.get("mu_val", 255), maintain_last=True + ) + + # unnorm action + if self.configs.get("norm_action", False): + from maniunicon.utils.vla_utils import unnoramalize_action_perdim + + if isinstance(action, tuple): + action = ( + unnoramalize_action_perdim( + action[0], np.array(self.configs["norm_min"]), np.array(self.configs["norm_max"]) + ), + action[1], + ) + else: + action = unnoramalize_action_perdim( + action, np.array(self.configs["norm_min"]), np.array(self.configs["norm_max"]), maintain_last=True + ) + + self.rollout_step_counter += 1 + if isinstance(action, torch.Tensor): + action = action.numpy() + + # convert rot from euler to quat + from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = action.shape + action = np.concatenate([euler_pose_to_quat(action[..., :-1].reshape(-1, A-1)).reshape(N, -1), action[..., -1:]], axis=-1) + print(f"step {self.rollout_step_counter} action {action}") + + return action diff --git a/maniunicon/customize/policy_model/spatialvla_model.py b/maniunicon/customize/policy_model/spatialvla_model.py new file mode 100644 index 0000000..ffd7778 --- /dev/null +++ b/maniunicon/customize/policy_model/spatialvla_model.py @@ -0,0 +1,178 @@ +from typing import Optional, Sequence, List +import os +import numpy as np +from transformers import AutoModel, AutoProcessor +from collections import deque +from PIL import Image +import torch +import cv2 as cv +import time + + +class ActionEnsembler: + def __init__(self, pred_action_horizon, action_ensemble_temp=0.0): + self.pred_action_horizon = pred_action_horizon + self.action_ensemble_temp = action_ensemble_temp + self.action_history = deque(maxlen=self.pred_action_horizon) + + def reset(self): + self.action_history.clear() + + def ensemble_action(self, cur_action): + self.action_history.append(cur_action) + num_actions = len(self.action_history) + if cur_action.ndim == 1: + curr_act_preds = np.stack(self.action_history) + else: + curr_act_preds = np.stack( + [pred_actions[i] for (i, pred_actions) in zip(range(num_actions - 1, -1, -1), self.action_history)] + ) + # if temp > 0, more recent predictions get exponentially *less* weight than older predictions + weights = np.exp(-self.action_ensemble_temp * np.arange(num_actions)) + weights = weights / weights.sum() + # compute the weighted average across all predictions for this timestep + cur_action = np.sum(weights[:, None] * curr_act_preds, axis=0) + + return cur_action + + +class SpatialVLAModel: + def __init__( + self, + device, + saved_model_path: str = "IPEC-COMMUNITY/spatialvla-4b-224-pt", + unnorm_key: str = None, + image_size: list[int] = [224, 224], + action_ensemble_temp: float = -0.8, + use_act_chunk: bool = False, + task_name: str = None + ) -> None: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + self.device = device + self.unnorm_key = unnorm_key + print(f"*** unnorm_key: {unnorm_key} ***") + # setup the task instruction that input to vla model + self.task_name = task_name + + self.processor = AutoProcessor.from_pretrained( + saved_model_path, trust_remote_code=True + ) + self.vla = ( + AutoModel.from_pretrained( + saved_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + .eval() + ) + self.vla.to(self.device) + + self.image_size = image_size + self.obs_horizon = (self.processor.num_obs_steps - 1) * self.processor.obs_delta + 1 + self.obs_interval = self.processor.obs_delta + self.pred_action_horizon = self.processor.action_chunk_size + self.image_history = deque(maxlen=self.obs_horizon) + self.use_act_chunk = use_act_chunk + + self.sticky_action_is_on = False + self.gripper_action_repeat = 0 + self.sticky_gripper_action = 0.0 + self.previous_gripper_action = None + + if self.use_act_chunk: + action_ensemble = False + else: + action_ensemble = True + self.action_ensemble = action_ensemble + self.action_ensemble_temp = action_ensemble_temp + + if self.action_ensemble: + self.action_ensembler = ActionEnsembler( + self.pred_action_horizon, self.action_ensemble_temp + ) + else: + self.action_ensembler = None + + self.rollout_step_counter = 0 + + def reset(self) -> None: + print("reset model now!!!!!!!!") + self.image_history.clear() + if self.action_ensemble: + self.action_ensembler.reset() + self.sticky_action_is_on = False + self.gripper_action_repeat = 0 + self.sticky_gripper_action = 0.0 + self.previous_gripper_action = None + + self.rollout_step_counter = 0 + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + image = cv.resize(image, tuple(self.image_size), interpolation=cv.INTER_AREA) + return image + + def _add_image_to_history(self, image: np.ndarray) -> None: + if len(self.image_history) == 0: + self.image_history.extend([image] * self.obs_horizon) + else: + self.image_history.append(image) + + def _obtain_image_history(self) -> List[Image.Image]: + image_history = list(self.image_history) + images = image_history[:: self.obs_interval] + images = [Image.fromarray(image).convert("RGB") for image in images] + return images + + def preprocess(self, obs): + obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() # (224, 224, 3) + # preprocess image + image = obs["image"]["camera_1"].numpy() + assert image.dtype == np.uint8 + image = self._resize_image(image) + self._add_image_to_history(image) + images: List[Image.Image] = self._obtain_image_history() + + return images + + def __call__(self, obs, **kwargs): + """ + Forward pass through the model. + :param obs: Observation input to the model. + :return: Model output. + """ + """Step function.""" + images = self.preprocess(obs) + prompt = self.task_name + + # predict action (7-dof; un-normalize for bridgev2) + start = time.time() + inputs = self.processor(images=images, text=prompt, unnorm_key=self.unnorm_key, return_tensors="pt", do_normalize=False) + with torch.no_grad(): + if hasattr(self.processor, "action_tokenizer"): + generation_outputs = self.vla.predict_action(inputs) + + raw_actions = self.processor.decode_actions( + generation_outputs=generation_outputs, + unnorm_key=self.unnorm_key, + )["actions"] + else: + raw_actions = self.vla.predict_action(**inputs)["actions"] + raw_actions = raw_actions.cpu().numpy() + # cal policy infer time + print("**** SpatialVLA inference time: ", time.time() - start) + + if self.action_ensemble: + print("ensemble action!!!") + raw_actions = self.action_ensembler.ensemble_action(raw_actions)[None] + + self.rollout_step_counter += 1 + if isinstance(raw_actions, torch.Tensor): + raw_actions = raw_actions.numpy() + + # convert rot from euler to quat + from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = raw_actions.shape + raw_actions = np.concatenate([euler_pose_to_quat(raw_actions[..., :-1].reshape(-1, A-1)).reshape(N, -1), raw_actions[..., -1:]], axis=-1) + print(f"step {self.rollout_step_counter} action {raw_actions}") + + return raw_actions diff --git a/maniunicon/policies/torch_model_vla.py b/maniunicon/policies/torch_model_vla.py new file mode 100644 index 0000000..0308a52 --- /dev/null +++ b/maniunicon/policies/torch_model_vla.py @@ -0,0 +1,336 @@ +import os +import time +import numpy as np +from typing import Callable, Any, Optional +import hydra +import torch +from multiprocessing.synchronize import Event +import traceback +from pynput import keyboard + +from maniunicon.utils.data import get_next_episode_dir +from maniunicon.utils.shared_memory.shared_storage import ( + SharedStorage, + MultiCameraData, +) +from maniunicon.core.policy import BasePolicy + + +def precise_wait(t_end: float, slack_time: float = 0.001, time_func=time.monotonic): + t_start = time_func() + t_wait = t_end - t_start + if t_wait > 0: + t_sleep = t_wait - slack_time + if t_sleep > 0: + time.sleep(t_sleep) + while time_func() < t_end: + pass + return + + +def compile(): + print("global compile setting..") + torch._dynamo.reset() + torch._dynamo.config.cache_size_limit = 256 + torch.set_float32_matmul_precision("high") + torch._dynamo.config.capture_scalar_outputs = True + torch._dynamo.config.capture_dynamic_output_shape_ops = True + + +class VLAPolicy(BasePolicy): + """Process of robot policy model.""" + + def __init__( + self, + shared_storage: SharedStorage, + reset_event: Event, + model: Callable, + obs_wrapper: Callable, + act_wrapper: Callable, + dt: float = 0.02, # policy future actions dt + name: str = "RoboVlmsPolicy", + steps_per_inference: int = 1, + infer_latency: float = 1.0 / 20, # seconds + frame_latency: float = 1.0 / 30, # seconds + use_real_time_chunking: bool = False, + record_dir: Optional[str] = None, + enable_recording: bool = False, + device: str = "cpu", + synchronized: bool = False, + **kwargs: Any, + ): + super().__init__( + shared_storage=shared_storage, + reset_event=reset_event, + name=name, + ) + self.model = model + self.obs_wrapper = obs_wrapper + self.act_wrapper = act_wrapper + self.dt = dt + self.obs_horizon = obs_wrapper.observation_horizon + self.steps_per_inference = steps_per_inference + self.frame_latency = frame_latency + self.use_real_time_chunking = use_real_time_chunking + self.enable_recording = enable_recording + self.infer_latency = infer_latency + self.device = device + self.synchronized = synchronized + self.activate = False + if self.use_real_time_chunking: + self.prev_raw_actions = None + self.prev_raw_timestamps = None + + self.record_dir = record_dir + self._recording_active = False + self._recording_key_pressed = False + self._current_episode_dir = None + self._listener = None + + def _on_press(self, key): + """Handle key press events.""" + try: + # Handle recording toggle on key press (not hold) + if ( + key == keyboard.KeyCode.from_char("r") + and not self._recording_key_pressed + ): + self._recording_key_pressed = True + self._recording_active = not self._recording_active + print(f"Recording: {'ON' if self._recording_active else 'OFF'}") + # Handle dropping current episode + elif key == keyboard.KeyCode.from_char("d") and self._recording_active: + self._recording_active = False + # First stop recording in shared storage + self.shared_storage.clear_record_dir() + # Then remove the directory if it exists + if self._current_episode_dir and os.path.exists( + self._current_episode_dir + ): + import shutil + + shutil.rmtree(self._current_episode_dir) + print( + f"Dropped episode - Removed directory: {self._current_episode_dir}" + ) + self._current_episode_dir = None + self.shared_storage.stop_record() + print("Recording: OFF") + except AttributeError: + pass + + def _on_release(self, key): + """Handle key release events.""" + try: + # Reset recording key press state + if key == keyboard.KeyCode.from_char("r"): + self._recording_key_pressed = False + except AttributeError: + pass + + def run(self): + """Main process loop.""" + try: + compile() # speedup model inference + self.model = hydra.utils.instantiate( + self.model, device=self.device, _recursive_=False + ) + self.obs_wrapper = hydra.utils.instantiate( + self.obs_wrapper, + shared_storage=self.shared_storage, + device=self.device, + ) + self.act_wrapper = hydra.utils.instantiate( + self.act_wrapper, + ) + + if self.enable_recording: + self._listener = keyboard.Listener( + on_press=self._on_press, on_release=self._on_release + ) + self._listener.start() + + while not self._recording_active: + print("Waiting for pressing 'r' to start recording...") + time.sleep(0.1) + + def _on_press(key): + if key == keyboard.KeyCode.from_char("g"): + print("Start policy inference!") + self.activate = True + + start_listener = keyboard.Listener(on_press=_on_press) + start_listener.start() + + self.model.reset() + start_delay = 1.0 + eval_t_start = time.time() + start_delay + t_start = time.monotonic() + start_delay + frame_latency = self.frame_latency + precise_wait(eval_t_start - frame_latency, time_func=time.time) + iter_idx = 0 + while self.shared_storage.is_running.value: + # calculate timing + if self.synchronized: + t_cycle_end = time.monotonic() + self.dt + iter_idx = int((t_cycle_end - t_start) / self.dt) + else: + t_cycle_end = ( + t_start + (iter_idx + self.steps_per_inference) * self.dt + ) + iter_idx += self.steps_per_inference + + if self.reset_event is not None and self.reset_event.is_set(): + self.shared_storage.read_all_action() + self.model.reset() + self.activate = False + precise_wait(t_cycle_end - frame_latency) + continue + + if not self.activate: + print("pressing 'g' to start policy inference...") + precise_wait(t_cycle_end - frame_latency) + continue + + # Synchronization logic: wait for robot to be ready before inference + if self.synchronized and ( + self.reset_event is None or not self.reset_event.is_set() + ): + # wait for robot to be ready + self.shared_storage.robot_ready.wait() + self.shared_storage.robot_ready.clear() + + if self.enable_recording: + # Handle recording state changes + if ( + self._recording_active + and not self.shared_storage.is_recording.value + ): + # Start recording + if self.record_dir is not None: + self._current_episode_dir = get_next_episode_dir( + self.record_dir + ) + self.shared_storage.set_record_dir( + self._current_episode_dir + ) + self.shared_storage.start_record( + start_time=time.time(), + dt=self.dt, + ) + print( + f"Recording started - Episode: {os.path.basename(self._current_episode_dir)}" + ) + else: + print( + "Recording directory not specified. Cannot start recording." + ) + self._recording_active = False + elif ( + not self._recording_active + and self.shared_storage.is_recording.value + ): + # Stop recording + self.shared_storage.stop_record() + if self._current_episode_dir: + print( + f"Recording stopped - Episode saved to: {self._current_episode_dir}" + ) + else: + print("Recording stopped") + self._current_episode_dir = None + self._recording_active = False + + state = self.shared_storage.read_state(k=self.obs_horizon) + camera: MultiCameraData = self.shared_storage.read_multi_camera( + k=self.obs_horizon + ) + obs_time = state.timestamp[-1].item() + + if state is None: + print("Robot is not ready, skipping policy") + precise_wait(t_cycle_end - frame_latency) + continue + if camera is None: + print("Camera is not ready, skipping policy") + precise_wait(t_cycle_end - frame_latency) + continue + + start = time.time() + # Get latest obs + obs: torch.Tensor = self.obs_wrapper(state=state, camera=camera) + + if self.use_real_time_chunking: + # if self.prev_raw_timestamps is not None: + # print("time diff", obs_time - self.prev_raw_timestamps[0]) + if self.prev_raw_actions is not None: + local_cond = np.zeros_like(self.prev_raw_actions) + overlap_idxes = np.where(self.prev_raw_timestamps > obs_time)[0] + num_overlap = overlap_idxes.shape[0] + hist_horizon = min(num_overlap, self.steps_per_inference) + local_cond[:, :, :hist_horizon] = self.prev_raw_actions[ + :, :, overlap_idxes[:hist_horizon] + ] + local_cond = torch.from_numpy(local_cond).to(self.device) + else: + local_cond = None + hist_horizon = 0 + actions, self.prev_raw_actions = self.act_wrapper( + self.model( + obs, + head_kwargs={ + "local_cond": local_cond, + "hist_horizon": hist_horizon, + }, + ), + timestamp=obs_time, + start_timestamp=eval_t_start, + ) + self.prev_raw_timestamps = ( + np.arange(self.prev_raw_actions.shape[2]) * self.dt + obs_time + ) + else: + actions = self.act_wrapper( + self.model(obs), + timestamp=obs_time, + start_timestamp=eval_t_start, + return_raw_actions=False, + ) + print( + "Obs wrapper + Policy Inference + Action wrapper: ", + time.time() - start, + ) + + # Check for errors + if self.shared_storage.error_state.value: + break + + if self.enable_recording and not self._recording_active: + precise_wait(t_cycle_end - frame_latency) + continue + + for action in actions: + self.shared_storage.write_action(action) + + # Synchronization logic: signal policy ready for robot execution + if self.synchronized and ( + self.reset_event is None or not self.reset_event.is_set() + ): + self.shared_storage.policy_ready.set() # Signal robot can execute actions + + if not self.synchronized: + precise_wait(t_cycle_end - frame_latency) + + except Exception as e: + print(f"Error in model: {e}") + traceback.print_exc() + self.shared_storage.error_state.value = True + # Reset synchronization state on error + if self.synchronized: + self.shared_storage.robot_ready.set() + self.shared_storage.policy_ready.clear() + + def stop(self): + """Stop the policy process.""" + self.shared_storage.is_running.value = False + self.join() diff --git a/maniunicon/utils/vla_utils.py b/maniunicon/utils/vla_utils.py new file mode 100644 index 0000000..9282b4e --- /dev/null +++ b/maniunicon/utils/vla_utils.py @@ -0,0 +1,28 @@ +import copy +import numpy as np +import scipy.spatial.transform as st + +def unnoramalize_action_perdim(action, action_min=-1, action_max=1, maintain_last=False): + last_val = copy.deepcopy(action[..., -1]) + res = 0.5 * (action + 1) * (action_max - action_min) + action_min + if maintain_last: + res[..., -1] = last_val + return res + +# Based on https://github.com/real-stanford/universal_manipulation_interface/blob/main/umi/common/pose_util.py +# convert from euler angle to rot +def euler_pose_to_pos_rot(euler_pose): + pos = euler_pose[..., :3] + rot = st.Rotation.from_euler('xyz', euler_pose[..., 3:], degrees=False) + return pos, rot + +def pos_rot_to_quat(pos, rot): + shape = pos.shape[:-1] + quat = np.zeros(shape + (7,), dtype=pos.dtype) + quat[...,:3] = pos + quat[...,3:] = rot.as_quat() + return quat + +# convert rot from euler angle to quat +def euler_pose_to_quat(euler_pose): + return pos_rot_to_quat(*euler_pose_to_pos_rot(euler_pose)) \ No newline at end of file From 819413595d1f31c0bc43746fe614e8643df8848d Mon Sep 17 00:00:00 2001 From: zbzhu99 Date: Mon, 3 Nov 2025 23:29:41 +0800 Subject: [PATCH 2/6] minor --- .../customize/obs_wrapper/robovlms_rgb_wrapper.py | 3 --- maniunicon/customize/policy_model/robovlms_model.py | 13 +++++++++---- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py index 87d25f7..2968be2 100644 --- a/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py +++ b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py @@ -23,9 +23,6 @@ def resize_image_sequence(images, target_size, interp=cv2.INTER_AREA): N, H, W, C = images.shape new_H, new_W = target_size - # Reshape to 2D array of images for faster processing - reshaped = images.reshape(-1, H, W, C) - # Preallocate output array output = np.empty((N, new_H, new_W, C), dtype=images.dtype) diff --git a/maniunicon/customize/policy_model/robovlms_model.py b/maniunicon/customize/policy_model/robovlms_model.py index 9cde5e7..86e2702 100644 --- a/maniunicon/customize/policy_model/robovlms_model.py +++ b/maniunicon/customize/policy_model/robovlms_model.py @@ -1,5 +1,6 @@ import json import os.path +import shutil from copy import deepcopy import torch from PIL import Image @@ -66,11 +67,15 @@ def init_config( import transformers package_dir = transformers.__path__[0] - os.system( - "cp /tools/modeling_kosmos2.py {}/models/kosmos2/modeling_kosmos2.py".format( - package_dir + robovlms_root = os.environ.get("ROBOVLMS_ROOT") + if not robovlms_root: + raise EnvironmentError( + "ROBOVLMS_ROOT environment variable must be set to the RoboVLMs installation directory." ) - ) + source_model_path = os.path.join(robovlms_root, "tools", "modeling_kosmos2.py") + target_model_path = os.path.join(package_dir, "models", "kosmos2", "modeling_kosmos2.py") + # Copy the custom kosmos model implementation into the transformers package when needed. + shutil.copy(source_model_path, target_model_path) if not self.debug: ckpt = torch.load(ckpt_path, map_location="cpu") From 9b147f5282f264255528bf408cbf05d589261df3 Mon Sep 17 00:00:00 2001 From: zbzhu99 Date: Mon, 3 Nov 2025 23:31:50 +0800 Subject: [PATCH 3/6] code format --- .../act_wrapper/robovlms_eepose_wrapper.py | 12 ++- .../act_wrapper/spatialvla_eepose_wrapper.py | 12 ++- .../obs_wrapper/robovlms_rgb_wrapper.py | 4 +- .../obs_wrapper/spatialvla_rgb_wrapper.py | 14 ++-- .../customize/policy_model/openvla_model.py | 1 + .../customize/policy_model/robovlms_model.py | 75 ++++++++++++++----- .../policy_model/spatialvla_model.py | 73 +++++++++++------- maniunicon/utils/vla_utils.py | 16 ++-- 8 files changed, 140 insertions(+), 67 deletions(-) diff --git a/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py b/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py index 62fd5f1..41a881b 100644 --- a/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py +++ b/maniunicon/customize/act_wrapper/robovlms_eepose_wrapper.py @@ -56,12 +56,18 @@ def __call__( raw_actions = actions.copy() if actions.ndim == 4: # action: (bs, ws, fwd, action_dim) remove the batch dimension and history actions - actions = actions[0, 0, self.hist_action : self.hist_action + self.action_horizon, :] + actions = actions[ + 0, 0, self.hist_action : self.hist_action + self.action_horizon, : + ] elif actions.ndim == 2: # action: (fwd, action_dim) - actions = actions[self.hist_action : self.hist_action + self.action_horizon, :] + actions = actions[ + self.hist_action : self.hist_action + self.action_horizon, : + ] else: - raise ValueError(f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array.") + raise ValueError( + f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array." + ) print(np.sum(is_new), " new actions, ", len(actions), " total actions") if np.sum(is_new) == 0: diff --git a/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py b/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py index aac8ef3..23ee608 100644 --- a/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py +++ b/maniunicon/customize/act_wrapper/spatialvla_eepose_wrapper.py @@ -56,12 +56,18 @@ def __call__( raw_actions = actions.copy() if actions.ndim == 4: # action: (bs, ws, fwd, action_dim) remove the batch dimension and history actions - actions = actions[0, 0, self.hist_action : self.hist_action + self.action_horizon, :] + actions = actions[ + 0, 0, self.hist_action : self.hist_action + self.action_horizon, : + ] elif actions.ndim == 2: # action: (fwd, action_dim) - actions = actions[self.hist_action : self.hist_action + self.action_horizon, :] + actions = actions[ + self.hist_action : self.hist_action + self.action_horizon, : + ] else: - raise ValueError(f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array.") + raise ValueError( + f"Unsupported action shape: {actions.shape}. Expected 2D or 4D array." + ) print(np.sum(is_new), " new actions, ", len(actions), " total actions") if np.sum(is_new) == 0: diff --git a/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py index 2968be2..a202a4c 100644 --- a/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py +++ b/maniunicon/customize/obs_wrapper/robovlms_rgb_wrapper.py @@ -28,9 +28,7 @@ def resize_image_sequence(images, target_size, interp=cv2.INTER_AREA): # Resize each image for i in range(N): - res = cv2.resize( - images[i], (new_W, new_H), interpolation=interp - ) + res = cv2.resize(images[i], (new_W, new_H), interpolation=interp) if C == 1: output[i] = res[:, :, np.newaxis] else: diff --git a/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py b/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py index 73625dd..35f634d 100644 --- a/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py +++ b/maniunicon/customize/obs_wrapper/spatialvla_rgb_wrapper.py @@ -27,26 +27,26 @@ def resize_image_sequence(images, target_size, resample=Image.BICUBIC): if isinstance(target_size, int): target_size = (target_size, target_size) new_H, new_W = target_size - + N, H, W, C = images.shape - + # Preallocate output array (always 3 channels for RGB) output = np.empty((N, new_H, new_W, 3), dtype=np.uint8) - + # Resize each image using PIL for i in range(N): # Convert numpy array to PIL Image img_pil = Image.fromarray(images[i]) - + # Convert to RGB (regardless of original number of channels) img_pil = img_pil.convert("RGB") - + # Resize image img_resized = img_pil.resize((new_W, new_H), resample=resample) - + # Convert back to numpy array and store in output output[i] = np.array(img_resized) - + return output diff --git a/maniunicon/customize/policy_model/openvla_model.py b/maniunicon/customize/policy_model/openvla_model.py index 97041e8..fb3c6e5 100644 --- a/maniunicon/customize/policy_model/openvla_model.py +++ b/maniunicon/customize/policy_model/openvla_model.py @@ -312,6 +312,7 @@ def __call__(self, obs, **kwargs): # OpenVLA outputs: [x, y, z, rx, ry, rz, gripper] # Need to convert to: [x, y, z, qx, qy, qz, qw, gripper] from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = raw_actions.shape if A == 7: # Euler format raw_actions = np.concatenate( diff --git a/maniunicon/customize/policy_model/robovlms_model.py b/maniunicon/customize/policy_model/robovlms_model.py index 86e2702..e65a0a7 100644 --- a/maniunicon/customize/policy_model/robovlms_model.py +++ b/maniunicon/customize/policy_model/robovlms_model.py @@ -24,6 +24,7 @@ fwd_decay_ratio = 1 + class RoboVlmsModel: # model option def __init__( @@ -35,30 +36,42 @@ def __init__( raw_calvin=True, debug=False, use_act_chunk=False, - task_name=None - ): + task_name=None, + ): # setup the task instruction that input to vla model self.task_name = task_name # Loading robovlms model configs assert config_path != None configs = load_config(config_path) self.model = BaseTrainer(configs=configs) - + # Get checkpoint path print("ckpt_path", ckpt_path) - from robovlms.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict + from robovlms.utils.zero_to_fp32 import ( + convert_zero_checkpoint_to_fp32_state_dict, + ) + # Handle DeepSpeed ckpt if os.path.isdir(ckpt_path): target_ckpt_path = ckpt_path.replace(".ckpt", ".pt") print(f"converting {ckpt_path} to {target_ckpt_path}") convert_zero_checkpoint_to_fp32_state_dict(ckpt_path, target_ckpt_path) ckpt_path = target_ckpt_path - - self.init_config(ckpt_path, configs, device, log_save_dir, raw_calvin, debug, use_act_chunk) + + self.init_config( + ckpt_path, configs, device, log_save_dir, raw_calvin, debug, use_act_chunk + ) # self.model.model.lm_head.window_size = 1 def init_config( - self, ckpt_path, configs, device, save_dir=None, raw_calvin=False, debug=False, use_act_chunk=False + self, + ckpt_path, + configs, + device, + save_dir=None, + raw_calvin=False, + debug=False, + use_act_chunk=False, ): ### load and convert checkpoint self.debug = debug @@ -72,8 +85,12 @@ def init_config( raise EnvironmentError( "ROBOVLMS_ROOT environment variable must be set to the RoboVLMs installation directory." ) - source_model_path = os.path.join(robovlms_root, "tools", "modeling_kosmos2.py") - target_model_path = os.path.join(package_dir, "models", "kosmos2", "modeling_kosmos2.py") + source_model_path = os.path.join( + robovlms_root, "tools", "modeling_kosmos2.py" + ) + target_model_path = os.path.join( + package_dir, "models", "kosmos2", "modeling_kosmos2.py" + ) # Copy the custom kosmos model implementation into the transformers package when needed. shutil.copy(source_model_path, target_model_path) @@ -257,12 +274,14 @@ def get_history(self, q: Queue, pad: Literal["zero", "first"] = "zero"): return queue_list, pad_mask def preprocess(self, obs, lang, mode="continuous"): - obs["image"]["camera_0"] = obs["image"]["camera_0"].squeeze().cpu().detach() # (224, 224, 3) + obs["image"]["camera_0"] = ( + obs["image"]["camera_0"].squeeze().cpu().detach() + ) # (224, 224, 3) obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() # preprocess static cam image image = obs["image"]["camera_1"].numpy() image = Image.fromarray(image) - image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) + image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) gripper_x = None if "camera_0" in obs["image"]: @@ -363,7 +382,9 @@ def __call__(self, obs, **kwargs): """ """Step function.""" input_dict = dict() - image_x, gripper_x, text_x, mask = self.preprocess(obs, self.task_name, self.action_space) + image_x, gripper_x, text_x, mask = self.preprocess( + obs, self.task_name, self.action_space + ) input_dict["rgb"] = image_x input_dict["hand_rgb"] = gripper_x @@ -379,14 +400,14 @@ def __call__(self, obs, **kwargs): action = self.policy.inference_step(input_dict)["action"] # print("original action from model: ", action) print("**** RoboVLMs inference time: ", time.time() - start) - + if self.action_space != "discrete": if action[0].ndim == action[1].ndim + 1: action = (action[0], action[1].unsqueeze(2)) action = torch.cat( [action[0], (torch.nn.functional.sigmoid(action[1]) > 0.5).float()], dim=-1, - ) + ) if isinstance(action, tuple): action = torch.cat([action[0], action[1]], dim=-1) @@ -417,8 +438,10 @@ def __call__(self, obs, **kwargs): action = self.ensemble_action(action) if isinstance(action, torch.Tensor): - action = action.squeeze() # (fwd, 7) for action chunk, (7,) for ensembled action - if not self.use_act_chunk: + action = ( + action.squeeze() + ) # (fwd, 7) for action chunk, (7,) for ensembled action + if not self.use_act_chunk: action = action.unsqueeze(0) if self.configs.get("use_mu_law", False): @@ -435,13 +458,18 @@ def __call__(self, obs, **kwargs): if isinstance(action, tuple): action = ( unnoramalize_action_perdim( - action[0], np.array(self.configs["norm_min"]), np.array(self.configs["norm_max"]) + action[0], + np.array(self.configs["norm_min"]), + np.array(self.configs["norm_max"]), ), action[1], ) else: action = unnoramalize_action_perdim( - action, np.array(self.configs["norm_min"]), np.array(self.configs["norm_max"]), maintain_last=True + action, + np.array(self.configs["norm_min"]), + np.array(self.configs["norm_max"]), + maintain_last=True, ) self.rollout_step_counter += 1 @@ -450,8 +478,15 @@ def __call__(self, obs, **kwargs): # convert rot from euler to quat from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = action.shape - action = np.concatenate([euler_pose_to_quat(action[..., :-1].reshape(-1, A-1)).reshape(N, -1), action[..., -1:]], axis=-1) + action = np.concatenate( + [ + euler_pose_to_quat(action[..., :-1].reshape(-1, A - 1)).reshape(N, -1), + action[..., -1:], + ], + axis=-1, + ) print(f"step {self.rollout_step_counter} action {action}") - + return action diff --git a/maniunicon/customize/policy_model/spatialvla_model.py b/maniunicon/customize/policy_model/spatialvla_model.py index ffd7778..d4b2555 100644 --- a/maniunicon/customize/policy_model/spatialvla_model.py +++ b/maniunicon/customize/policy_model/spatialvla_model.py @@ -25,7 +25,12 @@ def ensemble_action(self, cur_action): curr_act_preds = np.stack(self.action_history) else: curr_act_preds = np.stack( - [pred_actions[i] for (i, pred_actions) in zip(range(num_actions - 1, -1, -1), self.action_history)] + [ + pred_actions[i] + for (i, pred_actions) in zip( + range(num_actions - 1, -1, -1), self.action_history + ) + ] ) # if temp > 0, more recent predictions get exponentially *less* weight than older predictions weights = np.exp(-self.action_ensemble_temp * np.arange(num_actions)) @@ -34,7 +39,7 @@ def ensemble_action(self, cur_action): cur_action = np.sum(weights[:, None] * curr_act_preds, axis=0) return cur_action - + class SpatialVLAModel: def __init__( @@ -45,7 +50,7 @@ def __init__( image_size: list[int] = [224, 224], action_ensemble_temp: float = -0.8, use_act_chunk: bool = False, - task_name: str = None + task_name: str = None, ) -> None: os.environ["TOKENIZERS_PARALLELISM"] = "false" self.device = device @@ -53,22 +58,21 @@ def __init__( print(f"*** unnorm_key: {unnorm_key} ***") # setup the task instruction that input to vla model self.task_name = task_name - + self.processor = AutoProcessor.from_pretrained( saved_model_path, trust_remote_code=True ) - self.vla = ( - AutoModel.from_pretrained( - saved_model_path, - torch_dtype=torch.bfloat16, - trust_remote_code=True, - ) - .eval() - ) + self.vla = AutoModel.from_pretrained( + saved_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ).eval() self.vla.to(self.device) self.image_size = image_size - self.obs_horizon = (self.processor.num_obs_steps - 1) * self.processor.obs_delta + 1 + self.obs_horizon = ( + self.processor.num_obs_steps - 1 + ) * self.processor.obs_delta + 1 self.obs_interval = self.processor.obs_delta self.pred_action_horizon = self.processor.action_chunk_size self.image_history = deque(maxlen=self.obs_horizon) @@ -92,9 +96,9 @@ def __init__( ) else: self.action_ensembler = None - + self.rollout_step_counter = 0 - + def reset(self) -> None: print("reset model now!!!!!!!!") self.image_history.clear() @@ -104,19 +108,19 @@ def reset(self) -> None: self.gripper_action_repeat = 0 self.sticky_gripper_action = 0.0 self.previous_gripper_action = None - + self.rollout_step_counter = 0 - + def _resize_image(self, image: np.ndarray) -> np.ndarray: image = cv.resize(image, tuple(self.image_size), interpolation=cv.INTER_AREA) return image - + def _add_image_to_history(self, image: np.ndarray) -> None: if len(self.image_history) == 0: self.image_history.extend([image] * self.obs_horizon) else: self.image_history.append(image) - + def _obtain_image_history(self) -> List[Image.Image]: image_history = list(self.image_history) images = image_history[:: self.obs_interval] @@ -124,7 +128,9 @@ def _obtain_image_history(self) -> List[Image.Image]: return images def preprocess(self, obs): - obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() # (224, 224, 3) + obs["image"]["camera_1"] = ( + obs["image"]["camera_1"].squeeze().cpu().detach() + ) # (224, 224, 3) # preprocess image image = obs["image"]["camera_1"].numpy() assert image.dtype == np.uint8 @@ -143,14 +149,20 @@ def __call__(self, obs, **kwargs): """Step function.""" images = self.preprocess(obs) prompt = self.task_name - + # predict action (7-dof; un-normalize for bridgev2) start = time.time() - inputs = self.processor(images=images, text=prompt, unnorm_key=self.unnorm_key, return_tensors="pt", do_normalize=False) + inputs = self.processor( + images=images, + text=prompt, + unnorm_key=self.unnorm_key, + return_tensors="pt", + do_normalize=False, + ) with torch.no_grad(): if hasattr(self.processor, "action_tokenizer"): generation_outputs = self.vla.predict_action(inputs) - + raw_actions = self.processor.decode_actions( generation_outputs=generation_outputs, unnorm_key=self.unnorm_key, @@ -160,7 +172,7 @@ def __call__(self, obs, **kwargs): raw_actions = raw_actions.cpu().numpy() # cal policy infer time print("**** SpatialVLA inference time: ", time.time() - start) - + if self.action_ensemble: print("ensemble action!!!") raw_actions = self.action_ensembler.ensemble_action(raw_actions)[None] @@ -171,8 +183,17 @@ def __call__(self, obs, **kwargs): # convert rot from euler to quat from maniunicon.utils.vla_utils import euler_pose_to_quat + N, A = raw_actions.shape - raw_actions = np.concatenate([euler_pose_to_quat(raw_actions[..., :-1].reshape(-1, A-1)).reshape(N, -1), raw_actions[..., -1:]], axis=-1) + raw_actions = np.concatenate( + [ + euler_pose_to_quat(raw_actions[..., :-1].reshape(-1, A - 1)).reshape( + N, -1 + ), + raw_actions[..., -1:], + ], + axis=-1, + ) print(f"step {self.rollout_step_counter} action {raw_actions}") - + return raw_actions diff --git a/maniunicon/utils/vla_utils.py b/maniunicon/utils/vla_utils.py index 9282b4e..bbaa6a4 100644 --- a/maniunicon/utils/vla_utils.py +++ b/maniunicon/utils/vla_utils.py @@ -2,27 +2,33 @@ import numpy as np import scipy.spatial.transform as st -def unnoramalize_action_perdim(action, action_min=-1, action_max=1, maintain_last=False): + +def unnoramalize_action_perdim( + action, action_min=-1, action_max=1, maintain_last=False +): last_val = copy.deepcopy(action[..., -1]) res = 0.5 * (action + 1) * (action_max - action_min) + action_min if maintain_last: res[..., -1] = last_val return res + # Based on https://github.com/real-stanford/universal_manipulation_interface/blob/main/umi/common/pose_util.py # convert from euler angle to rot def euler_pose_to_pos_rot(euler_pose): pos = euler_pose[..., :3] - rot = st.Rotation.from_euler('xyz', euler_pose[..., 3:], degrees=False) + rot = st.Rotation.from_euler("xyz", euler_pose[..., 3:], degrees=False) return pos, rot + def pos_rot_to_quat(pos, rot): shape = pos.shape[:-1] quat = np.zeros(shape + (7,), dtype=pos.dtype) - quat[...,:3] = pos - quat[...,3:] = rot.as_quat() + quat[..., :3] = pos + quat[..., 3:] = rot.as_quat() return quat + # convert rot from euler angle to quat def euler_pose_to_quat(euler_pose): - return pos_rot_to_quat(*euler_pose_to_pos_rot(euler_pose)) \ No newline at end of file + return pos_rot_to_quat(*euler_pose_to_pos_rot(euler_pose)) From 4bddf3908a54e8834b4c634324455578c993e597 Mon Sep 17 00:00:00 2001 From: Trav1slaflame Date: Sun, 24 May 2026 23:02:45 +0800 Subject: [PATCH 4/6] added FALCON (From Spatial to Actions) & FALCON-PCD model support. --- configs/policy/falcon.yaml | 33 ++ configs/policy/falcon_pcd.yaml | 34 ++ .../obs_wrapper/falcon_pcd_wrapper.py | 177 +++++++ .../customize/policy_model/falcon_model.py | 494 ++++++++++++++++++ 4 files changed, 738 insertions(+) create mode 100644 configs/policy/falcon.yaml create mode 100644 configs/policy/falcon_pcd.yaml create mode 100644 maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py create mode 100644 maniunicon/customize/policy_model/falcon_model.py diff --git a/configs/policy/falcon.yaml b/configs/policy/falcon.yaml new file mode 100644 index 0000000..91f055e --- /dev/null +++ b/configs/policy/falcon.yaml @@ -0,0 +1,33 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 3 # expected openloop steps +use_real_time_chunking: false +enable_recording: false +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.falcon_model.FalconModel + ckpt_path: null # path to falcon model ckpt + config_path: null # path to falcon model config file + device: ${device} + log_save_dir: null # path to save ckpt loading msg, defult is null + use_act_chunk: true + task_name: "pick up anything" # task instruction that input to vla model + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.robovlms_rgb_wrapper.RoboVlmsImageWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + num_cams: 2 + observation_horizon: 1 # should be the same as the window size in vla model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.robovlms_eepose_wrapper.RoboVlmsEEPoseWrapper + hist_action: 0 # number of historical actions to consider + action_horizon: 5 # should be the same as the action chunk size in vla model + control_mode: cartesian + dt: ${policy.dt} \ No newline at end of file diff --git a/configs/policy/falcon_pcd.yaml b/configs/policy/falcon_pcd.yaml new file mode 100644 index 0000000..e5243de --- /dev/null +++ b/configs/policy/falcon_pcd.yaml @@ -0,0 +1,34 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 3 # expected openloop steps +use_real_time_chunking: false +enable_recording: false +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.falcon_model.FalconModel + ckpt_path: null # path to falcon model ckpt + config_path: null # path to falcon model config file + device: ${device} + log_save_dir: null # path to save ckpt loading msg, defult is null + use_act_chunk: true + task_name: "pick up anything" # task instruction that input to vla model + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.falcon_pcd_wrapper.FalconPCDWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + pcd_num_points: 2048 + num_cams: 2 + observation_horizon: 1 # should be the same as the window size in vla model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.robovlms_eepose_wrapper.RoboVlmsEEPoseWrapper + hist_action: 0 # number of historical actions to consider + action_horizon: 5 # should be the same as the action chunk size in vla model + control_mode: cartesian + dt: ${policy.dt} \ No newline at end of file diff --git a/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py b/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py new file mode 100644 index 0000000..5414482 --- /dev/null +++ b/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py @@ -0,0 +1,177 @@ +import torch +import torch.nn.functional as F +import numpy as np +import cv2 +import time +from pathlib import Path + +from maniunicon.utils.misc import dict_apply, convert_to_torch +from maniunicon.utils.pcd_utils import ( + add_gaussian_noise, + randomly_drop_point, + sample_pcd_data, +) +from maniunicon.utils.pcd_utils import ( + create_pointcloud_from_depth_batch, + create_pointcloud_from_rgbd_batch, + BOUND, + pcd_filter_bound_torch, + uniform_sampling_torch, + capture_pcd, + capture_depth, +) +from maniunicon.utils.math_utils import quat_from_matrix +from scipy.spatial.transform import Rotation + + +def resize_image_sequence(images, target_size, interp=cv2.INTER_AREA): + """ + Resize an image sequence using OpenCV + + Args: + images: numpy array of shape (N, H, W, C) where: + N = number of images + H = height + W = width + C = channels + target_size: tuple of (height, width) + + Returns: + resized images array of shape (N, new_H, new_W, C) + """ + N, H, W, C = images.shape + new_H, new_W = target_size + + # Reshape to 2D array of images for faster processing + reshaped = images.reshape(-1, H, W, C) + + # Preallocate output array + output = np.empty((N, new_H, new_W, C), dtype=images.dtype) + + # Resize each image + for i in range(N): + res = cv2.resize( + images[i], (new_W, new_H), interpolation=interp + ) + if C == 1: + output[i] = res[:, :, np.newaxis] + else: + output[i] = res + + return output + + +class FalconPCDWrapper: + def __init__( + self, + camera_config, + state_keys, + pcd_num_points=2048, + num_cams=1, + shared_storage=None, + device="cpu", + observation_horizon=1, + ): + self.camera_config = camera_config + self.state_keys = state_keys + self.num_cams = num_cams + self.pcd_num_points = pcd_num_points + self.shared_storage = shared_storage + self.device = device + + def __call__(self, state, camera): + """ + Wraps the observation for the Falcon-PCD model. + :return: Wrapped observation tensor. + """ + assert state is not None, "State cannot be None" + assert camera is not None, "Camera cannot be None" + + state_tensor = ( + torch.from_numpy( + np.concatenate( + [getattr(state, key) for key in self.state_keys], axis=-1 + ) + ) + .to(self.device) + .float() + ) + + colors = resize_image_sequence( + camera.colors.reshape(-1, *camera.colors.shape[2:]), (224, 224) + ) + colors = colors.reshape( + camera.colors.shape[0], + camera.colors.shape[1], + *colors.shape[1:], + ) + + images_tensor = {} + for cam_idx in range(colors.shape[1]): + images_tensor[f"camera_{cam_idx}"] = torch.from_numpy( + colors[:, cam_idx] + ).to(self.device) + + pcd_tensor = self.create_pcds_from_camera(camera, self.device) + + obs_tensor = { + "state": state_tensor, + "image": images_tensor, + "pointcloud": pcd_tensor + } + + return obs_tensor + + def create_pcds_from_camera(self, camera, device): + """ + Create point clouds from camera data. + :param camera: Camera object containing RGB, depth, intrinsics and extrinsics. + :return: Processed point clouds. + """ + CAMERA_IDX = 1 + + transforms = camera.transforms.reshape(-1, 4, 4)[CAMERA_IDX] + position = transforms[:3, 3] # (3) + # Extract rotation matrix and convert to quaternion (w, x, y, z) + rotation_matrix = transforms[:3, :3] + quat_xyzw = Rotation.from_matrix(rotation_matrix).as_quat() + # Convert to (w, x, y, z) format for create_pointcloud_from_depth_batch + orientation = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) + position = convert_to_torch( + position, dtype=torch.float32, device=device + ) + orientation = convert_to_torch( + orientation, dtype=torch.float32, device=device + ) + + intrinsics = convert_to_torch( + camera.intrs.reshape(-1, 4), dtype=torch.float32, device=device + ) # n_cams, 4 + intrinsics = intrinsics[CAMERA_IDX] + # Extract fx, fy, cx, cy from intrinsics + fx, fy, cx, cy = intrinsics[0], intrinsics[1], intrinsics[2], intrinsics[3] + # Build (3,3) intrinsic matrices for camera_1 + intrinsic_matrix = torch.zeros(3, 3) + intrinsic_matrix[0, 0] = fx + intrinsic_matrix[1, 1] = fy + intrinsic_matrix[0, 2] = cx + intrinsic_matrix[1, 2] = cy + intrinsic_matrix[2, 2] = 1.0 + + depths = convert_to_torch( + camera.depths.reshape(-1, camera.depths.shape[2], camera.depths.shape[3]), + dtype=torch.float32, + device=device, + ) # n_cams, H, W + depth = depths[CAMERA_IDX] + + pcd_tensor = create_pointcloud_from_depth_batch( + intrinsic_matrix=intrinsic_matrix, + depth=depth, + keep_invalid=False, + position=position, + orientation=orientation, + device=device + ) + + return pcd_tensor \ No newline at end of file diff --git a/maniunicon/customize/policy_model/falcon_model.py b/maniunicon/customize/policy_model/falcon_model.py new file mode 100644 index 0000000..39ec450 --- /dev/null +++ b/maniunicon/customize/policy_model/falcon_model.py @@ -0,0 +1,494 @@ +import json +import os.path +import shutil +from copy import deepcopy +import torch +from PIL import Image +from typing import Literal +import numpy as np +import functools +import time + +# Import FALCON components from the installed package +from falcon.train.base_trainer import BaseTrainer +from falcon.utils.model_utils import build_tokenizer +from falcon.data.data_utils import get_text_function +from falcon.data.data_utils import ( + preprocess_image, + get_prompt_builder, + tcp_to_world_frame, +) +from falcon.utils.config_utils import load_config +from queue import Queue +from falcon.model.policy_head.action_tokenizer import ActionTokenizer + +fwd_decay_ratio = 1 +# --------------------------------------------- +# Global configuration for pcd filtering: maximum distance from centroid (meters) +# Points farther than this from their frame centroid are considered noise. +# --------------------------------------------- +MAX_DISTANCE = 0.6 + +class FalconModel: + # model option + def __init__( + self, + ckpt_path, + config_path, + device, + log_save_dir=None, + raw_calvin=True, + debug=False, + use_act_chunk=False, + task_name=None, + ): + # setup the task instruction that input to vla model + self.task_name = task_name + # Loading falcon model configs + assert config_path != None + configs = load_config(config_path) + self.model = BaseTrainer(configs=configs) + + # Get checkpoint path + print("ckpt_path", ckpt_path) + from falcon.utils.zero_to_fp32 import ( + convert_zero_checkpoint_to_fp32_state_dict, + ) + + # Handle DeepSpeed ckpt + if os.path.isdir(ckpt_path): + target_ckpt_path = ckpt_path.replace(".ckpt", ".pt") + print(f"converting {ckpt_path} to {target_ckpt_path}") + convert_zero_checkpoint_to_fp32_state_dict(ckpt_path, target_ckpt_path) + ckpt_path = target_ckpt_path + + self.init_config( + ckpt_path, configs, device, log_save_dir, raw_calvin, debug, use_act_chunk + ) + # self.model.model.lm_head.window_size = 1 + + def init_config( + self, + ckpt_path, + configs, + device, + save_dir=None, + raw_calvin=False, + debug=False, + use_act_chunk=False, + ): + ### load and convert checkpoint + self.debug = debug + self.use_act_chunk = use_act_chunk + if configs["model"] == "kosmos": + import transformers + + package_dir = transformers.__path__[0] + falcon_root = os.environ.get("FALCON_ROOT") + if not falcon_root: + raise EnvironmentError( + "FALCON_ROOT environment variable must be set to the FALCON installation directory." + ) + source_model_path = os.path.join( + falcon_root, "tools", "modeling_kosmos2.py" + ) + target_model_path = os.path.join( + package_dir, "models", "kosmos2", "modeling_kosmos2.py" + ) + # Copy the custom kosmos model implementation into the transformers package when needed. + shutil.copy(source_model_path, target_model_path) + + if not self.debug: + ckpt = torch.load(ckpt_path, map_location="cpu") + if "state_dict" in ckpt: + new_state_dict = ckpt["state_dict"] + elif "model_state_dict" in ckpt: + new_state_dict = ckpt["model_state_dict"] + else: + raise KeyError("no checkpoint dict in the loaded pretrain parameters") + + new_state_dict = self.convert_old_state_dict(new_state_dict) + msg = self.model.load_state_dict(new_state_dict, strict=False) + print(f"FALCON CKPT Loaded \n {msg}") + del new_state_dict + + ckpt_dir = os.path.dirname(ckpt_path) + ckpt_name = os.path.basename(ckpt_path) + save_dir = ckpt_dir if save_dir is None else save_dir + load_info_path = os.path.join(save_dir, f"{ckpt_name}_loading_msg.json") + if os.path.exists(load_info_path): + os.system(f"rm {load_info_path}") + with open(load_info_path, "w") as f: + _info = { + "missing_keys": msg.missing_keys, + "unexpected_keys": msg.unexpected_keys, + } + json.dump(_info, f, indent=2) + print(f"Model loading msg is updated to: {load_info_path}") + + self.configs = configs + + dtype = torch.float32 + if self.configs["trainer"]["precision"] == "bf16": + dtype = torch.bfloat16 + elif self.configs["trainer"]["precision"] == "fp16": + dtype = torch.float16 + self.dtype = dtype + self.act_head_configs = self.configs["act_head"] + self.raw_calvin = raw_calvin + self.tcp_rel = self.configs.get("tcp_rel", False) + + print(f"raw action: {self.raw_calvin}") + + self.device = device + self.policy = self.model + self.policy = self.policy.to(self.dtype) + # self.policy = self.policy.float() + self.policy.to(self.device) + self.policy.eval() + + if not hasattr(self.policy.model, "lm_head"): + self.policy.model.lm_head = self.policy.model.act_head + + self.tokenizer = build_tokenizer(self.configs["tokenizer"]) + + self.window_size = configs["window_size"] + self.fwd_pred_next_n = configs["fwd_pred_next_n"] + self.act_step = self.fwd_pred_next_n + 1 + self.seq_len = self.configs["seq_len"] + self.use_hand_rgb = self.configs["use_hand_rgb"] + + if hasattr(self, "policy_setup"): + data_mix = "bridge" if self.policy_setup == "widowx_bridge" else "rt_1" + configs["train_dataset"]["data_mix"] = data_mix + configs["val_dataset"]["data_mix"] = data_mix + + image_preprocess = self.model.model.image_processor + self.image_preprocess = functools.partial( + preprocess_image, + image_processor=image_preprocess, + model_type=configs["model"], + ) + + self.text_preprocess = get_text_function( + self.model.model.tokenizer, configs["model"] + ) + + self.action_space = self.configs["act_head"].get("action_space", "continuous") + if self.action_space == "discrete": + self.action_tokenizer = ActionTokenizer( + self.tokenizer, + bins=self.act_head_configs["n_bin"], + min_action=self.act_head_configs["min_action"], + max_action=self.act_head_configs["max_action"], + ) + + print(f"Evaluating checkpoint {ckpt_path}") + + self.rgb_list = [] + self.hand_rgb_list = [] + self.action_hist_list = [] + self.rollout_step_counter = 0 + + self.vision_queue = Queue(maxsize=self.window_size) + self.vision_gripper_queue = Queue(maxsize=self.window_size) + self.action_queue = Queue(maxsize=self.window_size - 1) + + def ensemble_action(self, action): + if action.ndim >= 3: + action = action.squeeze() + + if action.ndim == 1: + action = action.unsqueeze(0) + + self.action_hist_list.append(action) + + act_cache = [] + max_len = 5 + while len(self.action_hist_list) > max_len: + self.action_hist_list.pop(0) + + idx = 0 + for act in self.action_hist_list[::-1]: + act_cache.append(act[idx]) + idx += 1 + + act_cache = torch.stack(act_cache, dim=0) + + weights = torch.tensor([fwd_decay_ratio ** i for i in range(len(act_cache))]) + weights = weights / weights.sum() + + weighted_act = (act_cache * weights.unsqueeze(1)).sum(dim=0) + + return weighted_act + + @staticmethod + def convert_old_state_dict(state_dict): + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("module."): + new_k = k.replace("module.", "") + else: + new_k = k + + if not new_k.startswith("model."): + new_k = "model." + new_k + + new_state_dict[new_k] = state_dict[k] + return new_state_dict + + def _get_default_calvin_config(self): + return { + "type": "DiskCalvinDataset", + "data_dir": "CALVIN/task_ABCD_D/val", + "c_act_scaler": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + } + + def add_element_to_queue(self, q: Queue, element): + while q.qsize() >= q.maxsize: + q.get() + q.put(element) + + def get_history(self, q: Queue, pad: Literal["zero", "first"] = "zero"): + queue_list = list(q.queue) + if len(queue_list) == 0: + return queue_list, None + history_type = self.configs["act_head"].get("history_type", "pre") + if history_type == "pre": + pad_len = 0 + else: + raise ValueError(f"Unsupported history type {history_type}") + element = queue_list[0] + if pad == "zero": + if isinstance(element, torch.Tensor): + element = torch.zeros_like(element) + elif isinstance(element, np.ndarray): + element = np.zeros_like(element) + else: + raise ValueError("This type is not supported") + queue_list = [element for _ in range(pad_len)] + queue_list + else: + if isinstance(element, torch.Tensor): + pad_list = [element.clone() for _ in range(pad_len)] + elif isinstance(element, np.ndarray): + pad_list = [deepcopy(element) for _ in range(pad_len)] + queue_list = pad_list + queue_list + pad_mask = np.ones(q.maxsize, dtype=bool) + pad_mask[:pad_len] = False + return queue_list, pad_mask + + def preprocess_esm(self, obs, lang, mode="continuous"): + obs["image"]["camera_0"] = obs["image"]["camera_0"].squeeze().cpu().detach() # (224, 224, 3) + obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() + # get inputs for esm + image_vggt = deepcopy(obs["image"]["camera_1"]).numpy() + # preprocess image + image = obs["image"]["camera_1"].numpy() + image = Image.fromarray(image) + image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) + + gripper_x = None + if "camera_0" in obs["image"]: + gripper = obs["image"]["camera_0"].numpy() + gripper = Image.fromarray(gripper) + gripper_x = self.image_preprocess([gripper]).unsqueeze(0) + gripper_x = gripper_x.to(self.device).to(self.dtype) + + # prepare inputs for esm + from falcon.model.policy_head.esm_utils.vggt.utils.load_fn import load_and_preprocess_images_square_new + # TODO: hardcode for param setup + vggt_target_size = 224 + image_vggt = Image.fromarray(image_vggt) + image_vggt_x, _ = load_and_preprocess_images_square_new([image_vggt], target_size=vggt_target_size) + image_vggt_x = image_vggt_x.unsqueeze(0).to(self.device).to(self.dtype) # (1, 1, 3, 224, 224) + + text_x, mask = self.text_preprocess([lang]) + + return ( + image_x.to(self.device).to(self.dtype), + gripper_x, + text_x.to(self.device), + mask.to(self.device), + image_vggt_x, + ) + + def preprocess_pcd(self, obs, lang, mode="continuous"): + obs["image"]["camera_0"] = obs["image"]["camera_0"].squeeze().cpu().detach() + obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() + # preprocess image + image = obs["image"]["camera_1"].numpy() + image = Image.fromarray(image) + image_x = self.image_preprocess([image]).unsqueeze(0) + + gripper_x = None + if "camera_0" in obs["image"]: + gripper = obs["image"]["camera_0"].numpy() + gripper = Image.fromarray(gripper) + gripper_x = self.image_preprocess([gripper]).unsqueeze(0) + gripper_x = gripper_x.to(self.device).to(self.dtype) + + # preprocess pcd + from falcon.eval.calvin.eval_utils import eval_filtered_sample_pcd + n_points = self.act_head_configs.get("pcd_n_points", 2048) + static_pcd_x = None + static_pcd = obs["pointcloud"].cpu().detach().numpy() # (N, 3) + static_pcd_x = eval_filtered_sample_pcd(static_pcd, n_points, MAX_DISTANCE).unsqueeze(0) + static_pcd_x = static_pcd_x.unsqueeze(0).to(self.device).to(self.dtype) + + text_x, mask = self.text_preprocess([lang]) + + return ( + image_x.to(self.device).to(self.dtype), + gripper_x, + text_x.to(self.device), + mask.to(self.device), + static_pcd_x, + ) + + def reset(self): + print("reset model now!!!!!!!!") + if hasattr(self.model.model, "lm_head"): + self.model.model.lm_head.hidden_state = None + self.model.model.lm_head.history_memory = [] + if hasattr(self.model.model, "act_head"): + self.model.model.act_head.hidden_state = None + self.model.model.act_head.history_memory = [] + self.model.model.act_head.esm_history_memory_rgb = [] + + self.rgb_list = [] + self.hand_rgb_list = [] + self.rollout_step_counter = 0 + self.action_hist_list = [] + + while not self.vision_queue.empty(): + self.vision_queue.get() + while not self.vision_gripper_queue.empty(): + self.vision_gripper_queue.get() + while not self.action_queue.empty(): + self.action_queue.get() + + def __call__(self, obs, **kwargs): + """ + Forward pass through the model. + :param obs: Observation input to the model. + :return: Model output. + """ + """Step function.""" + input_dict = dict() + + if self.act_head_configs["type"] in {"FCDecoder_ESM", "LSTMDecoder_ESM"}: + image_x, gripper_x, text_x, mask, image_esm_x = self.preprocess_esm(obs, self.task_name, self.action_space) + + input_dict["rgb"] = image_x + input_dict["hand_rgb"] = gripper_x + input_dict["text"] = text_x + input_dict["text_mask"] = mask + input_dict["rgb_vggt"] = image_esm_x + else: + image_x, gripper_x, text_x, mask, static_pcd_x = self.preprocess_pcd(obs, self.task_name, self.action_space) + + input_dict["rgb"] = image_x + input_dict["hand_rgb"] = gripper_x + input_dict["text"] = text_x + input_dict["text_mask"] = mask + input_dict["static_pcd"] = static_pcd_x + + if self.action_space == "discrete": + input_dict["instr_and_action_ids"] = text_x + input_dict["instr_and_action_mask"] = mask + + start = time.time() + with torch.no_grad(): + action = self.policy.inference_step(input_dict)["action"] + # print("original action from model: ", action) + print("**** FALCON inference time: ", time.time() - start) + + if self.action_space != "discrete": + if action[0].ndim == action[1].ndim + 1: + action = (action[0], action[1].unsqueeze(2)) + action = torch.cat( + [action[0], (torch.nn.functional.sigmoid(action[1]) > 0.5).float()], + dim=-1, + ) + + if isinstance(action, tuple): + action = torch.cat([action[0], action[1]], dim=-1) + + if isinstance(action, np.ndarray): + action = torch.from_numpy(action) + + if action.ndim == 2: + action = action.unsqueeze(1) + + if action.ndim == 3: + action = action.unsqueeze(1) + + action = action.detach().cpu() + + if self.tcp_rel: + robot_obs = ( + torch.from_numpy(obs["robot_obs"]) + .unsqueeze(0) + .unsqueeze(0) + .unsqueeze(0) + .repeat(1, self.window_size, self.fwd_pred_next_n, 1) + ) + action = tcp_to_world_frame(action, robot_obs) + + if not self.use_act_chunk: + print("ensemble action!!!") + action = self.ensemble_action(action) + + if isinstance(action, torch.Tensor): + action = ( + action.squeeze() + ) # (fwd, 7) for action chunk, (7,) for ensembled action + if not self.use_act_chunk: + action = action.unsqueeze(0) + + if self.configs.get("use_mu_law", False): + from falcon.data.data_utils import inverse_mu_law_companding + + action = inverse_mu_law_companding( + action, self.configs.get("mu_val", 255), maintain_last=True + ) + + # unnorm action + if self.configs.get("norm_action", False): + from maniunicon.utils.vla_utils import unnoramalize_action_perdim + + if isinstance(action, tuple): + action = ( + unnoramalize_action_perdim( + action[0], + np.array(self.configs["norm_min"]), + np.array(self.configs["norm_max"]), + ), + action[1], + ) + else: + action = unnoramalize_action_perdim( + action, + np.array(self.configs["norm_min"]), + np.array(self.configs["norm_max"]), + maintain_last=True, + ) + + self.rollout_step_counter += 1 + if isinstance(action, torch.Tensor): + action = action.numpy() + + # convert rot from euler to quat + from maniunicon.utils.vla_utils import euler_pose_to_quat + + N, A = action.shape + action = np.concatenate( + [ + euler_pose_to_quat(action[..., :-1].reshape(-1, A - 1)).reshape(N, -1), + action[..., -1:], + ], + axis=-1, + ) + print(f"step {self.rollout_step_counter} action {action}") + + return action From 3a531029f2970c9440fac45d6d55a039d47b15e3 Mon Sep 17 00:00:00 2001 From: Trav1slaflame Date: Sun, 24 May 2026 23:27:59 +0800 Subject: [PATCH 5/6] Apply black formatting for falcon_pcd_wrapper.py & falcon_model.py --- .../obs_wrapper/falcon_pcd_wrapper.py | 24 ++++------- .../customize/policy_model/falcon_model.py | 41 ++++++++++++++----- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py b/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py index 5414482..dd51a7e 100644 --- a/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py +++ b/maniunicon/customize/obs_wrapper/falcon_pcd_wrapper.py @@ -50,9 +50,7 @@ def resize_image_sequence(images, target_size, interp=cv2.INTER_AREA): # Resize each image for i in range(N): - res = cv2.resize( - images[i], (new_W, new_H), interpolation=interp - ) + res = cv2.resize(images[i], (new_W, new_H), interpolation=interp) if C == 1: output[i] = res[:, :, np.newaxis] else: @@ -117,7 +115,7 @@ def __call__(self, state, camera): obs_tensor = { "state": state_tensor, "image": images_tensor, - "pointcloud": pcd_tensor + "pointcloud": pcd_tensor, } return obs_tensor @@ -129,7 +127,7 @@ def create_pcds_from_camera(self, camera, device): :return: Processed point clouds. """ CAMERA_IDX = 1 - + transforms = camera.transforms.reshape(-1, 4, 4)[CAMERA_IDX] position = transforms[:3, 3] # (3) # Extract rotation matrix and convert to quaternion (w, x, y, z) @@ -137,13 +135,9 @@ def create_pcds_from_camera(self, camera, device): quat_xyzw = Rotation.from_matrix(rotation_matrix).as_quat() # Convert to (w, x, y, z) format for create_pointcloud_from_depth_batch orientation = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) - position = convert_to_torch( - position, dtype=torch.float32, device=device - ) - orientation = convert_to_torch( - orientation, dtype=torch.float32, device=device - ) - + position = convert_to_torch(position, dtype=torch.float32, device=device) + orientation = convert_to_torch(orientation, dtype=torch.float32, device=device) + intrinsics = convert_to_torch( camera.intrs.reshape(-1, 4), dtype=torch.float32, device=device ) # n_cams, 4 @@ -164,14 +158,14 @@ def create_pcds_from_camera(self, camera, device): device=device, ) # n_cams, H, W depth = depths[CAMERA_IDX] - + pcd_tensor = create_pointcloud_from_depth_batch( intrinsic_matrix=intrinsic_matrix, depth=depth, keep_invalid=False, position=position, orientation=orientation, - device=device + device=device, ) - return pcd_tensor \ No newline at end of file + return pcd_tensor diff --git a/maniunicon/customize/policy_model/falcon_model.py b/maniunicon/customize/policy_model/falcon_model.py index 39ec450..0bd0e5a 100644 --- a/maniunicon/customize/policy_model/falcon_model.py +++ b/maniunicon/customize/policy_model/falcon_model.py @@ -29,6 +29,7 @@ # --------------------------------------------- MAX_DISTANCE = 0.6 + class FalconModel: # model option def __init__( @@ -215,7 +216,7 @@ def ensemble_action(self, action): act_cache = torch.stack(act_cache, dim=0) - weights = torch.tensor([fwd_decay_ratio ** i for i in range(len(act_cache))]) + weights = torch.tensor([fwd_decay_ratio**i for i in range(len(act_cache))]) weights = weights / weights.sum() weighted_act = (act_cache * weights.unsqueeze(1)).sum(dim=0) @@ -278,14 +279,18 @@ def get_history(self, q: Queue, pad: Literal["zero", "first"] = "zero"): return queue_list, pad_mask def preprocess_esm(self, obs, lang, mode="continuous"): - obs["image"]["camera_0"] = obs["image"]["camera_0"].squeeze().cpu().detach() # (224, 224, 3) - obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() + obs["image"]["camera_0"] = ( + obs["image"]["camera_0"].squeeze().cpu().detach() + ) # (224, 224, 3) + obs["image"]["camera_1"] = ( + obs["image"]["camera_1"].squeeze().cpu().detach() + ) # get inputs for esm image_vggt = deepcopy(obs["image"]["camera_1"]).numpy() # preprocess image image = obs["image"]["camera_1"].numpy() image = Image.fromarray(image) - image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) + image_x = self.image_preprocess([image]).unsqueeze(0) # (1, 1, 3, 224, 224) gripper_x = None if "camera_0" in obs["image"]: @@ -295,12 +300,19 @@ def preprocess_esm(self, obs, lang, mode="continuous"): gripper_x = gripper_x.to(self.device).to(self.dtype) # prepare inputs for esm - from falcon.model.policy_head.esm_utils.vggt.utils.load_fn import load_and_preprocess_images_square_new + from falcon.model.policy_head.esm_utils.vggt.utils.load_fn import ( + load_and_preprocess_images_square_new, + ) + # TODO: hardcode for param setup vggt_target_size = 224 image_vggt = Image.fromarray(image_vggt) - image_vggt_x, _ = load_and_preprocess_images_square_new([image_vggt], target_size=vggt_target_size) - image_vggt_x = image_vggt_x.unsqueeze(0).to(self.device).to(self.dtype) # (1, 1, 3, 224, 224) + image_vggt_x, _ = load_and_preprocess_images_square_new( + [image_vggt], target_size=vggt_target_size + ) + image_vggt_x = ( + image_vggt_x.unsqueeze(0).to(self.device).to(self.dtype) + ) # (1, 1, 3, 224, 224) text_x, mask = self.text_preprocess([lang]) @@ -329,10 +341,13 @@ def preprocess_pcd(self, obs, lang, mode="continuous"): # preprocess pcd from falcon.eval.calvin.eval_utils import eval_filtered_sample_pcd + n_points = self.act_head_configs.get("pcd_n_points", 2048) static_pcd_x = None - static_pcd = obs["pointcloud"].cpu().detach().numpy() # (N, 3) - static_pcd_x = eval_filtered_sample_pcd(static_pcd, n_points, MAX_DISTANCE).unsqueeze(0) + static_pcd = obs["pointcloud"].cpu().detach().numpy() # (N, 3) + static_pcd_x = eval_filtered_sample_pcd( + static_pcd, n_points, MAX_DISTANCE + ).unsqueeze(0) static_pcd_x = static_pcd_x.unsqueeze(0).to(self.device).to(self.dtype) text_x, mask = self.text_preprocess([lang]) @@ -377,7 +392,9 @@ def __call__(self, obs, **kwargs): input_dict = dict() if self.act_head_configs["type"] in {"FCDecoder_ESM", "LSTMDecoder_ESM"}: - image_x, gripper_x, text_x, mask, image_esm_x = self.preprocess_esm(obs, self.task_name, self.action_space) + image_x, gripper_x, text_x, mask, image_esm_x = self.preprocess_esm( + obs, self.task_name, self.action_space + ) input_dict["rgb"] = image_x input_dict["hand_rgb"] = gripper_x @@ -385,7 +402,9 @@ def __call__(self, obs, **kwargs): input_dict["text_mask"] = mask input_dict["rgb_vggt"] = image_esm_x else: - image_x, gripper_x, text_x, mask, static_pcd_x = self.preprocess_pcd(obs, self.task_name, self.action_space) + image_x, gripper_x, text_x, mask, static_pcd_x = self.preprocess_pcd( + obs, self.task_name, self.action_space + ) input_dict["rgb"] = image_x input_dict["hand_rgb"] = gripper_x From f02f970f5945e8aeeab5f4f85369a90cdfd0d22e Mon Sep 17 00:00:00 2001 From: Trav1slaflame Date: Sun, 24 May 2026 23:30:57 +0800 Subject: [PATCH 6/6] minor fix to qualify the black format check :( --- maniunicon/customize/policy_model/falcon_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/maniunicon/customize/policy_model/falcon_model.py b/maniunicon/customize/policy_model/falcon_model.py index 0bd0e5a..83a4796 100644 --- a/maniunicon/customize/policy_model/falcon_model.py +++ b/maniunicon/customize/policy_model/falcon_model.py @@ -282,9 +282,7 @@ def preprocess_esm(self, obs, lang, mode="continuous"): obs["image"]["camera_0"] = ( obs["image"]["camera_0"].squeeze().cpu().detach() ) # (224, 224, 3) - obs["image"]["camera_1"] = ( - obs["image"]["camera_1"].squeeze().cpu().detach() - ) + obs["image"]["camera_1"] = obs["image"]["camera_1"].squeeze().cpu().detach() # get inputs for esm image_vggt = deepcopy(obs["image"]["camera_1"]).numpy() # preprocess image