from __future__ import annotations import inspect import sys from pathlib import Path from typing import Any import cv2 import numpy as np from XPolicyLab.model_template import ModelTemplate from XPolicyLab.utils.process_data import ( decode_image_bit, get_robot_action_dim_info, pack_robot_state, unpack_robot_state, ) _POLICY_DIR = Path(__file__).resolve().parent _MOTUS_ROOT = _POLICY_DIR / "motus" / "inference" / "robotwin" / "Motus" _CHECKPOINTS_DIR = _POLICY_DIR / "checkpoints" _DEFAULT_WAN_PATH = "/mnt/xspark-data/xspark_shared/model_weights/Wan2.2-TI2V-5B" _DEFAULT_VLM_PATH = "/mnt/xspark-data/xspark_shared/model_weights/Qwen3-VL-2B-Instruct" if str(_MOTUS_ROOT) not in sys.path: sys.path.insert(0, str(_MOTUS_ROOT)) from deploy_policy import MotusPolicy, get_model as get_motus_model, reset_model def extract_image(observation: dict[str, Any], candidate_names: list[str]) -> Any: vision = observation.get("vision", {}) for candidate_name in candidate_names: if candidate_name not in vision: continue image = vision[candidate_name] if isinstance(image, dict): for image_key in ("color", "rgb"): if image_key in image: return image[image_key] else: return image images = observation.get("images", {}) for candidate_name in candidate_names: if candidate_name in images: return images[candidate_name] raise KeyError(f"Could not find any image for candidates: {candidate_names}") def decode_compressed_image(image_buffer: np.ndarray) -> np.ndarray: return decode_image_bit(image_buffer) def ensure_hwc_uint8(image: Any) -> np.ndarray: if isinstance(image, (bytes, bytearray, memoryview)): image = decode_compressed_image(np.frombuffer(bytes(image), dtype=np.uint8)) image = np.asarray(image) if image.ndim == 1 and image.dtype == np.uint8: image = decode_compressed_image(image) if image.ndim != 3: raise ValueError(f"Expected image ndim=3, got shape {image.shape}") if np.issubdtype(image.dtype, np.floating): image = np.clip(image, 0.0, 1.0) image = (image * 255.0).astype(np.uint8) elif image.dtype != np.uint8: image = image.astype(np.uint8) if image.shape[-1] in (1, 3): return image if image.shape[0] in (1, 3): return np.transpose(image, (1, 2, 0)) raise ValueError(f"Unsupported image shape: {image.shape}") def _normalize_prompt_value(value: Any) -> str | None: if value is None: return None if isinstance(value, bytes): value = value.decode("utf-8", errors="ignore") elif isinstance(value, np.ndarray) and value.ndim == 0: value = value.item() elif isinstance(value, np.generic): value = value.item() if isinstance(value, (list, tuple)): for item in value: normalized = _normalize_prompt_value(item) if normalized is not None: return normalized return None if isinstance(value, str): value = value.strip() return value or None return str(value) def resolve_prompt(observation: dict[str, Any], default_prompt: str | None) -> str: for key in ("prompt", "instruction", "task", "language_instruction"): prompt = _normalize_prompt_value(observation.get(key)) if prompt is not None: return prompt fallback = _normalize_prompt_value(default_prompt) if fallback is None: raise ValueError("No valid prompt found in observation or model config.") return fallback def encode_obs(observation: dict[str, Any], action_type: str, robot_action_dim_info: dict[str, Any]) -> dict[str, Any]: if "observation" in observation and "joint_action" in observation: return observation if robot_action_dim_info is None: raise ValueError("env_cfg_type is required when encoding raw environment observations for Motus.") head = ensure_hwc_uint8(extract_image(observation, ["cam_high", "cam_head", "head_camera", "top_camera"])) head = cv2.resize(head, (320, 240), interpolation=cv2.INTER_AREA) left = ensure_hwc_uint8(extract_image(observation, ["cam_left_wrist", "left_camera", "left_wrist", "wrist_left"])) right = ensure_hwc_uint8(extract_image(observation, ["cam_right_wrist", "right_camera", "right_wrist", "wrist_right"])) state = pack_robot_state(observation, action_type, robot_action_dim_info, source_type="obs").astype(np.float32) return { "observation": { "head_camera": {"rgb": head}, "left_camera": {"rgb": left}, "right_camera": {"rgb": right}, }, "joint_action": {"vector": state}, } def _as_int_or_default(value: Any, default: int) -> int: if value is None or value == "": return default return int(value) def patch_motus_runtime_config(model_cfg: dict[str, Any]) -> None: # The RoboDojo checkpoint was trained with the LeRobot config, while the # bundled RobotWin inference config has a shorter action sequence. common_overrides = { "global_downsample_rate": _as_int_or_default(model_cfg.get("global_downsample_rate"), 1), "video_action_freq_ratio": _as_int_or_default(model_cfg.get("video_action_freq_ratio"), 6), } if not getattr(MotusPolicy, "_xpolicylab_config_patched", False): original_create_model_config = MotusPolicy._create_model_config def _create_model_config_with_xpolicylab_overrides(self): common_cfg = self.config_dict.setdefault("common", {}) for key, value in getattr(MotusPolicy, "_xpolicylab_common_overrides", {}).items(): common_cfg[key] = value return original_create_model_config(self) MotusPolicy._xpolicylab_original_create_model_config = original_create_model_config MotusPolicy._create_model_config = _create_model_config_with_xpolicylab_overrides MotusPolicy._xpolicylab_config_patched = True MotusPolicy._xpolicylab_common_overrides = common_overrides def normalize_motus_image_layout(encoded_observation: dict[str, Any]) -> None: obs_data = encoded_observation.get("observation", {}) head_camera = obs_data.get("head_camera", {}) head_image = head_camera.get("rgb") if head_image is None: return head_image = ensure_hwc_uint8(head_image) if head_image.shape[:2] != (240, 320): head_image = cv2.resize(head_image, (320, 240), interpolation=cv2.INTER_AREA) head_camera["rgb"] = head_image def patch_motus_qwen_rope_index(policy: Any) -> None: vlm_model = getattr(getattr(getattr(policy, "model", None), "vlm_model", None), "model", None) if vlm_model is None or getattr(vlm_model, "_xpolicylab_rope_index_patched", False): return original_get_rope_index = vlm_model.get_rope_index if "mm_token_type_ids" in inspect.signature(original_get_rope_index).parameters: return def _get_rope_index_compat(*args, **kwargs): kwargs.pop("mm_token_type_ids", None) return original_get_rope_index(*args, **kwargs) vlm_model.get_rope_index = _get_rope_index_compat vlm_model._xpolicylab_rope_index_patched = True def _resolve_path(value: Any, base_dir: Path = _POLICY_DIR) -> Path | None: if value is None or value == "": return None path = Path(str(value)).expanduser() if not path.is_absolute(): path = (base_dir / path).resolve() return path.resolve() def resolve_motus_checkpoint(model_cfg: dict[str, Any]) -> str: for key in ("ckpt_setting", "checkpoint_path", "model_path"): explicit_path = _resolve_path(model_cfg.get(key)) if explicit_path is not None: return str(explicit_path) ckpt_name = model_cfg.get("ckpt_name") if ckpt_name: raw_ckpt_name = Path(str(ckpt_name)).expanduser() if raw_ckpt_name.is_absolute() or "/" in str(ckpt_name): return str(_resolve_path(ckpt_name)) tuple_keys = ("dataset_name", "ckpt_name", "env_cfg_type", "expert_data_num", "action_type", "seed") if all(model_cfg.get(key) is not None for key in tuple_keys): checkpoint_setting = "-".join(str(model_cfg[key]) for key in tuple_keys) tuple_path = (_CHECKPOINTS_DIR / checkpoint_setting).resolve() if tuple_path.exists(): return str(tuple_path) return str((_CHECKPOINTS_DIR / str(ckpt_name)).resolve()) raise ValueError("ckpt_name, ckpt_setting, checkpoint_path, or model_path is required for Motus.") class Model(ModelTemplate): def __init__(self, model_cfg: dict[str, Any]): self.model_cfg = dict(model_cfg) self.task_name = self.model_cfg.get("task_name", "default_task") self.action_type = self.model_cfg.get("action_type", "joint") if self.action_type != "joint": raise ValueError("Motus in XPolicyLab currently supports only action_type='joint'.") env_cfg = self.model_cfg.get("env_cfg") or self.model_cfg.get("env_cfg_type") self.robot_action_dim_info = get_robot_action_dim_info(env_cfg) if env_cfg is not None else None self.default_prompt = self.model_cfg.get("prompt", self.task_name) self._latest_env_idx_list: list[int] = [0] self.observation_window: list[dict[str, Any]] | None = None ckpt_setting = resolve_motus_checkpoint(self.model_cfg) model_args = dict(self.model_cfg) model_args["ckpt_setting"] = ckpt_setting model_args["wan_path"] = str(_resolve_path(model_args.get("wan_path")) or Path(_DEFAULT_WAN_PATH)) model_args["vlm_path"] = str(_resolve_path(model_args.get("vlm_path")) or Path(_DEFAULT_VLM_PATH)) model_args["prompt"] = model_args.get("prompt") or self.default_prompt patch_motus_runtime_config(model_args) self.policy = get_motus_model(model_args) patch_motus_qwen_rope_index(self.policy) self.model = self.policy def update_obs(self, obs): self.update_obs_batch([obs]) def update_obs_batch(self, obs_list): self._latest_env_idx_list = [obs.get("env_idx", index) for index, obs in enumerate(obs_list)] self.observation_window = [ { "observation": encode_obs(obs, self.action_type, self.robot_action_dim_info), "instruction": resolve_prompt(obs, self.default_prompt), } for obs in obs_list ] def get_action(self, **kwargs): action_list = self.get_action_batch(env_idx_list=[self._latest_env_idx_list[0]], **kwargs) return action_list[0] def get_action_batch(self, env_idx_list=None, **kwargs): if self.observation_window is None: raise AssertionError("update_obs or update_obs_batch first!") env_idx_list = env_idx_list or self._latest_env_idx_list action_list = [] for batch_index, _ in enumerate(env_idx_list): cached = self.observation_window[batch_index] normalize_motus_image_layout(cached["observation"]) self.policy.set_instruction(cached["instruction"]) self.policy.update_obs(cached["observation"]) action_chunk = self.policy.get_action() action_list.append( unpack_robot_state( action_chunk, self.action_type, self.robot_action_dim_info, source_type="obs", ) ) return action_list def reset(self): self.observation_window = None self._latest_env_idx_list = [0] reset_model(self.policy) def reset_obsrvationwindows(self): self.reset()