Update Wall-X to 1.1.0 (#104)

wall_x/_vendor/harrix/serving/__init__.py

@@ -0,0 +1,3 @@
+from .websocket_policy_server import WebsocketPolicyServer, BasePolicy
+
+__all__ = ["WebsocketPolicyServer", "BasePolicy"]

wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py

@@ -0,0 +1 @@
+"""Wall-X inference helpers vendored for harrix serving."""

wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py

@@ -0,0 +1,476 @@
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from typing import Optional, List
+from dataclasses import dataclass, field
+import numpy as np
+import torch
+import wall_x._vendor.x2robot_utils.geometry as data_utils
+from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger
+
+dof_dims = {
+    "left_ee_cartesian_pos": 3,
+    "left_ee_cartesian_pos_relative": 3,
+    "left_ee_rotation": 3,
+    "left_ee_rotation_relative": 3,
+    "left_ee_rotation_6D": 6,
+    "left_ee_rotation_6D_relative": 6,
+    "left_arm_joint_pos": 7,  # 6 joint + 1 gripper
+    "left_gripper": 1,
+    "left_gripper_cur": 1,
+    "left_arm_joint_cur": 7,  # 1 -> 7
+    "right_ee_cartesian_pos": 3,
+    "right_ee_cartesian_pos_relative": 3,
+    "right_ee_rotation": 3,
+    "right_ee_rotation_relative": 3,
+    "right_ee_rotation_6D": 6,
+    "right_ee_rotation_6D_relative": 6,
+    "right_arm_joint_pos": 7,  # 6 joint + 1 gripper
+    "right_gripper": 1,
+    "right_gripper_cur": 1,
+    "right_arm_joint_cur": 7,  # 1 -> 7
+    "head_actions": 2,
+    "height": 1,
+    "car_pose": 3,
+    "velocity_decomposed": 3,
+    "velocity_decomposed_odom": 3,
+    "head_rotation": 2,  # match ex001
+    "left_joint": 6,
+    "left_joint_gripper": 1,
+    "right_joint": 6,
+    "right_joint_gripper": 1,
+    "left_rotation_quat": 4,
+    "right_rotation_quat": 4,
+    "left_quaternion": 4,
+    "right_quaternion": 4,
+    "left_wrench_ext_local_force": 3,
+    "left_wrench_ext_local_torque": 3,
+    "right_wrench_ext_local_force": 3,
+    "right_wrench_ext_local_torque": 3,
+    "left_wrench_ext_local_force_from_joint": 3,
+    "left_wrench_ext_local_torque_from_joint": 3,
+    "right_wrench_ext_local_force_from_joint": 3,
+    "right_wrench_ext_local_torque_from_joint": 3,
+    "left_wrench_ext_world_force": 3,
+    "left_wrench_ext_world_torque": 3,
+    "left_wrench_ext_world_force_from_joint": 3,
+    "left_wrench_ext_world_torque_from_joint": 3,
+    "right_wrench_ext_world_force": 3,
+    "right_wrench_ext_world_torque": 3,
+    "right_wrench_ext_world_force_from_joint": 3,
+    "right_wrench_ext_world_torque_from_joint": 3,
+    "left_arm_joint_dev": 7,
+    "right_arm_joint_dev": 7,
+}
+
+
+class ComputedDict(dict):
+    """Dict that registers compute rules and auto-computes None values on get"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._compute_rules = {}  # key -> compute_function
+
+    def register_compute_rule(self, key, compute_func):
+        """
+        Register a compute rule
+
+        Args:
+            key: Key to compute
+            compute_func: Callable taking self, returns computed value
+        """
+        self._compute_rules[key] = compute_func
+
+    def get(self, key, default=None):
+        """Override get to support auto-compute"""
+        value = super().get(key, default)
+
+        # If value is None and a rule exists, try to compute
+        if value is None and key in self._compute_rules:
+            try:
+                computed_value = self._compute_rules[key](self)
+                if computed_value is not None:
+                    # Cache computed value
+                    self[key] = computed_value
+                    return computed_value
+            except Exception:
+                pass  # On failure return None or default
+
+        return value if value is not None else default
+
+    def __getitem__(self, key):
+        """Override [] to support auto-compute"""
+        value = super().get(key, None)
+
+        # If value is None and a rule exists, try to compute
+        if value is None and key in self._compute_rules:
+            try:
+                computed_value = self._compute_rules[key](self)
+                if computed_value is not None:
+                    # Cache computed value
+                    self[key] = computed_value
+                    return computed_value
+            except Exception:
+                pass  # On failure raise KeyError or return None
+
+        if key in self:
+            return super().__getitem__(key)
+        raise KeyError(key)
+
+
+@dataclass
+class RobotStateActionData:
+    config: InferConfig = None
+    data: ComputedDict = field(
+        default_factory=lambda: ComputedDict(
+            {
+                # State (formerly pose) - state_ prefix
+                "state_left_ee_cartesian_pos": None,  # (1, 3)
+                "state_left_ee_rotation": None,  # (1, 3)
+                "state_left_ee_rotation_6D": None,
+                "state_left_arm_joint_pos": None,
+                "state_left_gripper": None,  # (1, 1)
+                "state_left_arm_joint_cur": None,
+                "state_left_gripper_cur": None,
+                "state_right_ee_cartesian_pos": None,  # (1, 3)
+                "state_right_ee_rotation": None,
+                "state_right_ee_rotation_6D": None,  # (1, 6)
+                "state_right_arm_joint_pos": None,
+                "state_right_gripper": None,
+                "state_right_gripper_cur": None,
+                "state_right_arm_joint_cur": None,  # (1, 1)
+                "state_head_actions": None,
+                "state_head_rotation": None,  # match ex001
+                "state_height": None,
+                "state_car_pose": None,
+                "state_velocity_decomposed": None,
+                "state_velocity_decomposed_odom": None,
+                # support joint control
+                "state_left_joint": None,
+                "state_left_joint_gripper": None,
+                "state_right_joint": None,
+                "state_right_joint_gripper": None,
+                # support quaternion control
+                "state_left_quaternion": None,
+                "state_right_quaternion": None,
+                "state_left_rotation_quat": None,
+                "state_right_rotation_quat": None,
+                # support wrench/force-torque observation
+                "state_left_wrench_ext_local_force": None,
+                "state_left_wrench_ext_local_torque": None,
+                "state_right_wrench_ext_local_force": None,
+                "state_right_wrench_ext_local_torque": None,
+                "state_left_wrench_ext_local_force_from_joint": None,
+                "state_left_wrench_ext_local_torque_from_joint": None,
+                "state_right_wrench_ext_local_force_from_joint": None,
+                "state_right_wrench_ext_local_torque_from_joint": None,
+                "state_left_wrench_ext_world_force": None,
+                "state_left_wrench_ext_world_torque": None,
+                "state_left_wrench_ext_world_force_from_joint": None,
+                "state_left_wrench_ext_world_torque_from_joint": None,
+                "state_right_wrench_ext_world_force": None,
+                "state_right_wrench_ext_world_torque": None,
+                "state_right_wrench_ext_world_force_from_joint": None,
+                "state_right_wrench_ext_world_torque_from_joint": None,
+                # support joint deviation
+                "state_left_arm_joint_dev": None,
+                "state_right_arm_joint_dev": None,
+                # Action - action_ prefix
+                "action_left_ee_cartesian_pos": None,
+                "action_left_ee_cartesian_pos_relative": None,
+                "action_left_ee_rotation": None,
+                "action_left_ee_rotation_relative": None,
+                "action_left_ee_rotation_6D": None,
+                "action_left_ee_rotation_6D_relative": None,
+                "action_left_gripper": None,
+                "action_left_arm_joint_pos": None,
+                "action_right_ee_cartesian_pos": None,
+                "action_right_ee_cartesian_pos_relative": None,
+                "action_right_ee_rotation": None,
+                "action_right_ee_rotation_relative": None,
+                "action_right_ee_rotation_6D": None,
+                "action_right_ee_rotation_6D_relative": None,
+                "action_right_gripper": None,
+                "action_right_arm_joint_pos": None,
+                "action_head_actions": None,
+                "action_head_rotation": None,  # match ex001
+                "action_height": None,
+                "action_car_pose": None,
+                "action_velocity_decomposed": None,
+                "action_velocity_decomposed_odom": None,
+                # support joint control
+                "action_left_joint": None,
+                "action_left_joint_gripper": None,
+                "action_right_joint": None,
+                "action_right_joint_gripper": None,
+                # support quaternion control
+                "action_left_quaternion": None,
+                "action_right_quaternion": None,
+                "action_left_rotation_quat": None,
+                "action_right_rotation_quat": None,
+            }
+        )
+    )
+    dof_mask: np.ndarray = None
+    logger = InferLogger.get_robot_logger("RobotStateActionData")
+
+    def __post_init__(self):
+        """Register compute rules"""
+        # State rules: euler angles -> 6D rotation
+        self.data.register_compute_rule(
+            "state_left_ee_rotation_6D",
+            lambda d: (
+                data_utils.euler_to_matrix_zyx_6d_nb(d["state_left_ee_rotation"])
+                if d.get("state_left_ee_rotation") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(
+            "state_right_ee_rotation_6D",
+            lambda d: (
+                data_utils.euler_to_matrix_zyx_6d_nb(d["state_right_ee_rotation"])
+                if d.get("state_right_ee_rotation") is not None
+                else None
+            ),
+        )
+
+        # Action rules: absolute position from relative + state
+        self.data.register_compute_rule(
+            "action_left_ee_cartesian_pos",
+            lambda d: (
+                d.get("state_left_ee_cartesian_pos")
+                + d.get("action_left_ee_cartesian_pos_relative")
+                if d.get("state_left_ee_cartesian_pos") is not None
+                and d.get("action_left_ee_cartesian_pos_relative") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(
+            "action_right_ee_cartesian_pos",
+            lambda d: (
+                d.get("state_right_ee_cartesian_pos")
+                + d.get("action_right_ee_cartesian_pos_relative")
+                if d.get("state_right_ee_cartesian_pos") is not None
+                and d.get("action_right_ee_cartesian_pos_relative") is not None
+                else None
+            ),
+        )
+
+        # Action rules: compute absolute rpy
+        self.data.register_compute_rule(  # delta rpy -> abs rpy
+            "action_left_ee_rotation",
+            lambda d: (
+                data_utils.compose_state_and_delta_to_abs_rpy(
+                    d["action_left_ee_rotation_relative"],
+                    d["state_left_ee_rotation"][0],
+                )
+                if d.get("action_left_ee_rotation_relative") is not None
+                and d.get("state_left_ee_rotation") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(  # abs 6D -> abs rpy
+            "action_left_ee_rotation",
+            lambda d: (
+                data_utils.so3_to_euler_zyx_batch_nb(d["action_left_ee_rotation_6D"])
+                if d.get("action_left_ee_rotation_6D") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(  # delta 6D -> abs 6D
+            "action_left_ee_rotation_6D",
+            lambda d: (
+                data_utils.compose_state_and_delta_to_abs_6d(
+                    d["action_left_ee_rotation_6D_relative"],
+                    d["state_left_ee_rotation_6D"][0],
+                )
+                if d.get("action_left_ee_rotation_6D_relative") is not None
+                and d.get("state_left_ee_rotation_6D") is not None
+                else None
+            ),
+        )
+
+        self.data.register_compute_rule(  # delta rpy -> abs rpy
+            "action_right_ee_rotation",
+            lambda d: (
+                data_utils.compose_state_and_delta_to_abs_rpy(
+                    d["action_right_ee_rotation_relative"],
+                    d["state_right_ee_rotation"][0],
+                )
+                if d.get("action_right_ee_rotation_relative") is not None
+                and d.get("state_right_ee_rotation") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(  # abs 6D -> abs rpy
+            "action_right_ee_rotation",
+            lambda d: (
+                data_utils.so3_to_euler_zyx_batch_nb(d["action_right_ee_rotation_6D"])
+                if d.get("action_right_ee_rotation_6D") is not None
+                else None
+            ),
+        )
+        self.data.register_compute_rule(  # delta 6D -> abs 6D
+            "action_right_ee_rotation_6D",
+            lambda d: (
+                data_utils.compose_state_and_delta_to_abs_6d(
+                    d["action_right_ee_rotation_6D_relative"],
+                    d["state_right_ee_rotation_6D"][0],
+                )
+                if d.get("action_right_ee_rotation_6D_relative") is not None
+                and d.get("state_right_ee_rotation_6D") is not None
+                else None
+            ),
+        )
+
+    def get_agent_pos(self, obs_action_keys=None):
+        if obs_action_keys is None:
+            obs_action_keys = self.config.train_config["agent_pos_config"].keys()
+
+        agent_pose_data = []
+        for key in obs_action_keys:
+            # action_padding is a virtual key used to pad agent_pos width.
+            if key == "action_padding":
+                dim = self.config.train_config["agent_pos_config"][key]
+                agent_pose_data.append(np.zeros((1, dim)))
+                continue
+
+            # Strip follow_ or master_ prefix
+            if key.startswith("follow_"):
+                key = key.replace("follow_", "")
+            elif key.startswith("master_"):
+                key = key.replace("master_", "")
+
+            # Add state_ prefix to access state data
+            state_key = f"state_{key}"
+
+            if state_key in self.data:
+                # get() auto-computes when value is None
+                value = self.data.get(state_key)
+                if value is None:
+                    # If still None after compute, use zeros
+                    agent_pose_data.append(np.zeros((1, dof_dims[key])))
+                else:
+                    agent_pose_data.append(value)
+            else:
+                raise ValueError(f"Key {state_key} not found in data")
+
+        agent_pose_data = np.concatenate(agent_pose_data, axis=1)[None]  # (1, 1, D)
+
+        return agent_pose_data
+
+    def get_agent_pos_mask(self, obs_action_keys=None):
+        if obs_action_keys is None:
+            obs_action_keys = self.config.train_config["agent_pos_config"].keys()
+
+        agent_pos_mask_data = []
+        for key in obs_action_keys:
+            # action_padding carries no information and should stay masked out.
+            if key == "action_padding":
+                dim = self.config.train_config["agent_pos_config"][key]
+                agent_pos_mask_data.append(np.zeros((1, dim)))
+                continue
+
+            # Strip follow_ or master_ prefix
+            if key.startswith("follow_"):
+                key = key.replace("follow_", "")
+            elif key.startswith("master_"):
+                key = key.replace("master_", "")
+
+            # Add state_ prefix to access state data
+            state_key = f"state_{key}"
+
+            if state_key in self.data:
+                # get() auto-computes when value is None
+                value = self.data.get(state_key)
+                if value is None:
+                    agent_pos_mask_data.append(np.zeros((1, dof_dims[key])))
+                else:
+                    agent_pos_mask_data.append(np.ones((1, value.shape[1])))
+            else:
+                raise ValueError(f"Key {state_key} not found in data")
+
+        return np.concatenate(agent_pos_mask_data, axis=1)[None]  # (1, 1, D)
+
+    def save_state_data_with_key(self, value, key, gt_dim=None):
+        # Strip follow_ or master_ prefix
+        key = key.replace("follow_", "")
+        key = key.replace("master_", "")
+
+        # if torch, convert to numpy
+        if isinstance(value, torch.Tensor):
+            value = value.detach().cpu().numpy()
+
+        if f"state_{key}" not in self.data:  # TODO: joint angle control
+            self.logger.warning(f"{key} is not a valid state key; not recorded")
+            return
+        gt_dim = dof_dims[key] if gt_dim is None else gt_dim
+
+        # Shape check: expected (1, D)
+        if value.shape == (1, gt_dim):
+            self.data[f"state_{key}"] = value
+        elif value.shape == (1, 1, gt_dim):
+            self.data[f"state_{key}"] = value[0]
+        elif value.shape == (gt_dim,):
+            self.data[f"state_{key}"] = value[None]
+        else:
+            raise ValueError(f"Value shape {value.shape} is not legal")
+
+    def save_action_data_with_key(self, value, key):
+        key = key.replace("follow_", "")
+        key = key.replace("master_", "")
+
+        if isinstance(value, torch.Tensor):
+            value = value.detach().cpu().numpy()
+
+        if value.shape == (dof_dims[key],):
+            self.data[f"action_{key}"] = value[None]
+        else:
+            self.data[f"action_{key}"] = value
+
+    def save_action_data(
+        self, predict_action, predict_action_keys: Optional[List[str]] = None
+    ):
+        if predict_action_keys is None:
+            predict_action_keys = getattr(
+                self.config.data_config, "predict_action_keys", None
+            )
+            if predict_action_keys is None:
+                try:
+                    predict_action_keys = self.config.data_config["predict_action_keys"]
+                except (KeyError, AttributeError):
+                    predict_action_keys = list(
+                        self.config.train_config["dof_config"].keys()
+                    )
+
+        if isinstance(predict_action, torch.Tensor):
+            predict_action = predict_action.detach().cpu().numpy()
+
+        if predict_action.ndim == 3:
+            predict_action = predict_action[0]
+
+        dof_start = 0
+        for action_key in predict_action_keys:
+            # action_padding is a virtual key that only advances dof_start; no data is written.
+            if action_key == "action_padding":
+                dof_dim = self.config.train_config["dof_config"]["action_padding"]
+                dof_start += dof_dim
+                continue
+
+            action_key = action_key.replace("follow_", "")
+            action_key = action_key.replace("master_", "")
+            dof_dim = dof_dims[action_key]
+            action_key = f"action_{action_key}"
+            self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim]
+            dof_start += dof_dim
+
+    # Convenience properties for compatibility
+    @property
+    def agent_pos(self):
+        return self.get_agent_pos()
+
+    @property
+    def agent_pos_mask(self):
+        return self.get_agent_pos_mask()
+
+    @property
+    def action(self):
+        pass  # TODO: support action access

wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py

@@ -0,0 +1,434 @@
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+import yaml
+from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS
+
+logger = logging.getLogger(__name__)
+
+
+def _env_bool(name: str, default: bool) -> bool:
+    v = os.environ.get(name)
+    if v is None:
+        return default
+    return v.strip().lower() in ("1", "true", "yes", "on")
+
+
+@dataclass
+class InferenceDataConfig:
+    """Minimal data config needed by online inference.
+
+    Serving does not build datasets, so it should not require closed-source
+    data backends just to read image resize settings from a checkpoint yaml.
+    """
+
+    resolution: dict[str, int] = field(default_factory=dict)
+    model_type: str = "qwen2_5"
+    max_pixels: int = MAX_PIXELS
+    min_pixels: int = MIN_PIXELS
+    image_factor: int = IMAGE_FACTOR
+    use_relative_action: bool = False
+    predict_action_keys: list[str] = field(default_factory=list)
+
+    def __getitem__(self, key: str) -> Any:
+        if not hasattr(self, key):
+            raise KeyError(key)
+        return getattr(self, key)
+
+
+class InferConfig:
+    def __init__(
+        self,
+        checkpoint_path: str | None = None,
+        train_config_path: str | None = None,
+        robot_host: str = "0.0.0.0",
+        robot_port: int = 41776,
+        robot_type: str = "desktop",  # ["desktop", "turtle", "ex001"]
+        robot_action_start_ratio: float = 0,  # Start ratio for trimming executed actions
+        robot_action_end_ratio: float = 0.8,  # End ratio for trimming executed actions
+        robot_action_interpolate_multiplier: int = 10,  # Action interpolation multiplier
+        robot_use_joint_angle_control: bool = False,  # Joint control (model must predict joints)
+        turtle_as_desktop: bool = False,  # Use turtle as desktop: fixed base/head motion, head camera, base height
+        action_horizon: int = 32,  # Set to the model's action horizon
+        action_dim: int | None = None,
+        ar_action_dim: int | None = None,
+        model_device: str = "cuda:0",
+        num_inference_timesteps: int = 10,
+        num_inference_steps: int | None = None,
+        cfg_scale: float | None = None,
+        seed: int | None = None,
+        norm_key: str = "x2_normal",  # ["x2_normal", "ex_normal"]
+        cam_names: list[str] | None = None,
+        camera_front_key: str = "camera_front",
+        camera_left_key: str = "camera_left",
+        camera_right_key: str = "camera_right",
+        default_instruction: str | None = None,
+        prompt_template: str | None = None,
+        qwen25_prompt_template: str | None = None,
+        prompt_priority_order: str | None = None,
+        save_video_dir: str = "./videos",
+        robot_id: str = "10000",
+        model_type: str = "wallx",  # ["wallx", "vga"]
+        smooth_action: bool = False,
+        smooth_gripper: bool = True,
+    ):
+        # Private attributes for paths
+        assert checkpoint_path is not None
+        from wall_x._vendor.harrix.utils.ckpt_load import resolve_checkpoint_dir
+
+        if not os.path.isdir(checkpoint_path) and not os.path.isfile(checkpoint_path):
+            raise FileNotFoundError(
+                f"Checkpoint path not found: {checkpoint_path!r}. "
+                "Serving requires a directory containing model.safetensors and "
+                "normalizer_action.pth / normalizer_propri.pth (or norm_stats.json)."
+            )
+        self._checkpoint_path = resolve_checkpoint_dir(checkpoint_path)
+        action_pth = os.path.join(self._checkpoint_path, "normalizer_action.pth")
+        propri_pth = os.path.join(self._checkpoint_path, "normalizer_propri.pth")
+        if os.path.exists(action_pth):
+            self.normalizer_action_path = action_pth
+        if os.path.exists(propri_pth):
+            self.normalizer_propri_path = propri_pth
+
+        # Other config attributes
+        self.robot_host = robot_host
+        self.robot_port = robot_port
+        self.robot_type = robot_type  # ["desktop", "turtle", "ex001"]
+        self.robot_action_start_ratio = robot_action_start_ratio
+        self.robot_action_end_ratio = robot_action_end_ratio
+        self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier
+        self.robot_use_joint_angle_control = (
+            robot_use_joint_angle_control  # Joint-angle control
+        )
+        self.turtle_as_desktop = turtle_as_desktop
+        self.robot_id = robot_id
+        self._action_horizon = (
+            action_horizon  # Default from train config flow action horizon
+        )
+        self._action_dim = action_dim  # Default from train config dof config
+        self.ar_action_dim = ar_action_dim  # Default from train config ar dof config
+        self.model_device = model_device
+        self.num_inference_timesteps = (
+            num_inference_timesteps  # flow matching related config
+        )
+        self.num_inference_steps = num_inference_steps
+        self.cfg_scale = cfg_scale
+        self.seed = seed
+        self.save_video_dir = save_video_dir
+        self.model_type = model_type  # ["wallx", "vga"]
+        # Initialize config objects
+        self.train_config: dict = {}
+        self.model_config = None
+        self.data_config = None
+        self.norm_key = norm_key
+        self.cam_names = cam_names or [
+            "face_view",
+            "left_wrist_view",
+            "right_wrist_view",
+        ]
+        self.camera_front_key = camera_front_key
+        self.camera_left_key = camera_left_key
+        self.camera_right_key = camera_right_key
+        self.default_instruction = default_instruction
+        self.prompt_template = prompt_template
+        self.qwen25_prompt_template = qwen25_prompt_template
+        self.prompt_priority_order = prompt_priority_order
+        self.smooth_action = _env_bool("WALLX_SMOOTH_ACTION", smooth_action)
+        self.smooth_gripper = _env_bool("WALLX_SMOOTH_GRIPPER", smooth_gripper)
+        # Load all configs
+        self._load_all_configs(train_config_path)
+        self._apply_cam_names_from_train_config(cam_names)
+
+    def _apply_cam_names_from_train_config(
+        self, cam_names: list[str] | None
+    ) -> None:
+        """Use train YAML camera mapping when CLI did not override cam_names."""
+        if cam_names is not None:
+            return
+        from wall_x._vendor.harrix.utils.train_config import (
+            resolve_cam_names_from_train_config,
+        )
+
+        resolved = resolve_cam_names_from_train_config(self.train_config)
+        if resolved:
+            self.cam_names = resolved
+            logger.info(
+                "[InferConfig] cam_names from train key_mappings: %s",
+                self.cam_names,
+            )
+
+    @property
+    def checkpoint_path(self) -> str | None:
+        return self._checkpoint_path
+
+    @checkpoint_path.setter
+    def checkpoint_path(self, value: str | None):
+        """Reload all configs when checkpoint_path is updated"""
+        if self._checkpoint_path != value:
+            self._checkpoint_path = value
+            self._load_all_configs()
+
+    @property
+    def action_horizon(self) -> int:
+        return self._action_horizon
+
+    @action_horizon.setter
+    def action_horizon(self, value: int):
+        self._action_horizon = value
+
+    @property
+    def action_dim(self) -> int | None:
+        return self._action_dim
+
+    @action_dim.setter
+    def action_dim(self, value: int | None):
+        self._action_dim = value
+
+    @property
+    def ar_action_dim(self) -> int | None:
+        return self._ar_action_dim
+
+    @ar_action_dim.setter
+    def ar_action_dim(self, value: int | None):
+        self._ar_action_dim = value
+
+    def _load_all_configs(self, train_config_path=None):
+        """Unified entry to load all configs"""
+        self._load_train_config(train_config_path)
+        if self.model_type != "vga":
+            self._load_model_config_for_wallx()
+        self._load_data_config()
+
+        # Update action_horizon and action_dim if needed
+        if self._action_horizon is None:
+            self._action_horizon = self.train_config.get("data", {}).get(
+                "action_horizon_flow", 32
+            )
+        assert self._action_horizon is not None and self._action_horizon > 0
+
+        if self._action_dim is None:
+            dof_config = (
+                self.train_config.get("dof_config")
+                or self.train_config.get("task", {}).get("dof_config")
+                or self.train_config.get("data", {}).get("dof_config", {})
+            )
+            self._action_dim = sum(dof_config.values())
+
+        if self._ar_action_dim is None:
+            ar_dof_config = self.train_config.get("ar_dof_config") or self.train_config.get(
+                "task", {}
+            ).get("ar_dof_config", {})
+            self._ar_action_dim = sum(ar_dof_config.values())
+
+    def _load_train_config(self, train_config_path):
+        if train_config_path is None:
+            for fname in ("config.yml", "config.yaml"):
+                candidate = os.path.join(self._checkpoint_path, fname)
+                if os.path.exists(candidate):
+                    train_config_path = candidate
+                    break
+            else:
+                raise FileNotFoundError(
+                    f"No config.yml/config.yaml found in {self._checkpoint_path}"
+                )
+        with open(train_config_path, "r") as f:
+            self.train_config = yaml.load(f, Loader=yaml.FullLoader)
+
+        ckpt_dir = self._checkpoint_path
+        preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json")
+        if os.path.exists(preprocessor_file):
+            logger.info(
+                "[LoadConfig] Found %s, override processor_path.",
+                preprocessor_file,
+            )
+            self.train_config["processor_path"] = ckpt_dir
+
+        tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json")
+        tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json")
+        if self.train_config.get(
+            "action_tokenizer_path", None
+        ) is not None and not os.path.exists(
+            self.train_config.get("action_tokenizer_path", None)
+        ):
+            if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file):
+                logger.info(
+                    "[LoadConfig] Found tokenizer files in %s, override action_tokenizer_path.",
+                    ckpt_dir,
+                )
+                self.train_config["action_tokenizer_path"] = ckpt_dir
+            else:
+                logger.warning("[LoadConfig] Cannot load action tokenizer! ")
+
+        from wall_x._vendor.harrix.utils.train_config import (
+            normalize_train_config_for_inference,
+            strip_action_tokenizer_fields,
+        )
+
+        self.train_config = normalize_train_config_for_inference(
+            self.train_config, train_config_path
+        )
+        # Flow serving does not load AR action tokenizers; crop embed_tokens instead.
+        strip_action_tokenizer_fields(self.train_config)
+        self._train_config_path = train_config_path
+
+    def _load_model_config_for_wallx(self):
+        model_type = self.train_config["model_type"]
+        if not os.path.isdir(self._checkpoint_path):
+            return
+        # For Qwen models
+        ckpt_config_path = os.path.join(self._checkpoint_path, "config.json")
+        resolved_cfg_path = None
+
+        if os.path.exists(ckpt_config_path):
+            # Prefer checkpoint config
+            resolved_cfg_path = ckpt_config_path
+            logger.info(
+                "[LoadModelConfig] Using checkpoint config.json: %s",
+                ckpt_config_path,
+            )
+        else:
+            # Fallback to original config path
+            fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None)
+            if fallback_cfg is not None:
+                resolved_cfg_path = fallback_cfg
+                logger.info(
+                    "[LoadModelConfig] Using fallback act config: %s",
+                    fallback_cfg,
+                )
+
+        if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)):
+            raise ValueError(
+                f"[LoadModelConfig] Cannot load model config! "
+                f"Checked:\n"
+                f" - Checkpoint config.json: {ckpt_config_path}\n"
+                f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}"
+            )
+
+        # Save back to config for consistency
+        self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path
+
+        from wall_x.trainer.adapters import resolve_adapter
+
+        adapter_cls = resolve_adapter(model_type)
+        ConfigClass = adapter_cls.config_class()
+
+        logger.info(
+            "[LoadModelConfig] Loading model config from: %s", resolved_cfg_path
+        )
+        if resolved_cfg_path.endswith(".json"):
+            self.model_config = ConfigClass.from_json_file(resolved_cfg_path)
+        else:
+            self.model_config = ConfigClass.from_pretrained(resolved_cfg_path)
+
+        self.model_config.update_model_config(self.train_config)
+
+        self.model_config._attn_implementation = "sdpa"
+        vision_attn = os.environ.get("WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2")
+        self.model_config.vision_config._attn_implementation = vision_attn
+        logger.info(
+            "[LoadModelConfig] vision _attn_implementation=%s (override via WALLX_VISION_ATTN_IMPLEMENTATION)",
+            vision_attn,
+        )
+
+        logger.info("[LoadModelConfig] Model config loaded and updated successfully.")
+
+    def _load_data_config(self):
+        # Prefer typed TrainConfig path (handles new 8-section schema where
+        # dof_config lives under task:). Fall back to legacy raw-dict path
+        # for old flat yamls. See trainer/adapters/base_adapter.py.
+        typed_dcfg = self._try_typed_data_config()
+        if typed_dcfg is not None:
+            self.data_config = typed_dcfg
+        elif self.model_type != "vga":
+            self.data_config = self._build_inference_data_config()
+        else:
+            # TEMPORARY: direct call to the private _set_data_backend.
+            # VGA inference still uses the legacy x2robot data config object.
+            # Wall-X online inference uses the lightweight InferenceDataConfig
+            # above so serving does not require excluded data backends.
+            from wall_x.data._registry import _set_data_backend
+
+            _set_data_backend(self.train_config.get("dataset_type", "x2robot_v1"))
+            from wall_x.model.vga.openloop_visualization import get_data_configs
+
+            dataload_config = get_data_configs(self.train_config.get("data", {}))
+            dataload_config["predict_action_keys"] = list(
+                dataload_config.get("dof_config", {}).keys()
+            )
+            self.data_config = dataload_config
+
+        self._ensure_predict_action_keys()
+
+    def _ensure_predict_action_keys(self) -> None:
+        dof_keys = list((self.train_config.get("dof_config") or {}).keys())
+        if not dof_keys:
+            return
+        if isinstance(self.data_config, InferenceDataConfig):
+            if not self.data_config.predict_action_keys:
+                self.data_config.predict_action_keys = dof_keys
+        elif isinstance(self.data_config, dict):
+            self.data_config.setdefault("predict_action_keys", dof_keys)
+
+    def _build_inference_data_config(self) -> InferenceDataConfig:
+        data = self.train_config.get("data", {}) or {}
+
+        def get(key: str, default: Any) -> Any:
+            return data.get(key, self.train_config.get(key, default))
+
+        resolution = get("resolution", {}) or {}
+        dof_config = self.train_config.get("dof_config") or {}
+        return InferenceDataConfig(
+            resolution=dict(resolution),
+            model_type=get(
+                "model_type", self.train_config.get("model_type", "qwen2_5")
+            ),
+            max_pixels=get("max_pixels", MAX_PIXELS),
+            min_pixels=get("min_pixels", MIN_PIXELS),
+            image_factor=get("image_factor", IMAGE_FACTOR),
+            use_relative_action=get("use_relative_action", False),
+            predict_action_keys=list(dof_config.keys()),
+        )
+
+    def _try_typed_data_config(self):
+        """Attempt to load TrainConfig and build X2RDataConfig via typed path.
+
+        Returns the X2RDataConfig on success, or None if the yaml is not in
+        TrainConfig schema (legacy flat yaml).
+        """
+        yml_path = getattr(self, "_train_config_path", None)
+        if yml_path is None:
+            for fname in ("config.yml", "config.yaml"):
+                candidate = os.path.join(self._checkpoint_path, fname)
+                if os.path.exists(candidate):
+                    yml_path = candidate
+                    break
+        if yml_path is None:
+            return None
+        try:
+            from wall_x.config.loader import load_config
+            from wall_x.trainer.adapters.base_adapter import load_trainer_data_config
+
+            typed_cfg = load_config(yml_path)
+        except Exception as e:
+            logger.warning(
+                "[InferConfig] Typed config load failed for %s: %s",
+                yml_path,
+                e,
+            )
+            return None
+        try:
+            return load_trainer_data_config(typed_cfg)
+        except Exception as e:
+            logger.warning(
+                "[InferConfig] Typed config load failed for %s: %s",
+                yml_path,
+                e,
+            )
+            return None
+
+
+if __name__ == "__main__":
+    config = InferConfig()
+    logger.info("%s", config.train_config)

wall_x/_vendor/harrix/serving/_wallx_infer/logger.py

@@ -0,0 +1,277 @@
+"""Layered logging helpers for inference components."""
+
+import logging
+import sys
+from typing import Optional
+from pathlib import Path
+from datetime import datetime
+
+try:
+    import colorlog
+
+    HAS_COLORLOG = True
+except ImportError:
+    HAS_COLORLOG = False
+    logging.getLogger(__name__).warning(
+        "colorlog not installed. Install with: pip install colorlog"
+    )
+
+
+class InferLogger:
+    """
+    Layered inference logging system
+    """
+
+    _loggers = {}
+    _initialized = False
+
+    # Layer identifiers
+    LEVEL_ENV = "ENV"
+    LEVEL_ROBOT = "ROBOT"
+    LEVEL_CONTROLLER = "CONTROLLER"
+    LEVEL_MODEL = "MODEL"
+    LEVEL_UTILS = "UTILS"
+
+    # Layer colors (terminal output)
+    LEVEL_COLORS = {
+        LEVEL_ENV: "cyan",
+        LEVEL_ROBOT: "green",
+        LEVEL_CONTROLLER: "yellow",
+        LEVEL_MODEL: "purple",  # colorlog uses 'purple' not 'magenta'
+        LEVEL_UTILS: "blue",
+    }
+
+    @classmethod
+    def setup(
+        cls,
+        log_level: str = "INFO",
+        log_dir: Optional[str] = None,
+        console_output: bool = True,
+        file_output: bool = True,
+        colorful: bool = True,
+    ):
+        """
+        Initialize the logging system
+
+        Args:
+            log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+            log_dir: Log file directory
+            console_output: Whether to log to console
+            file_output: Whether to log to file
+            colorful: Colored console output (requires colorlog)
+        """
+        if cls._initialized:
+            return
+
+        cls.log_level = getattr(logging, log_level.upper())
+        cls.console_output = console_output
+        cls.file_output = file_output
+        cls.colorful = colorful and HAS_COLORLOG
+
+        # Create log directory
+        if file_output and log_dir:
+            cls.log_dir = Path(log_dir)
+            cls.log_dir.mkdir(parents=True, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            cls.log_file = cls.log_dir / f"infer_{timestamp}.log"
+        else:
+            cls.log_file = None
+
+        cls._initialized = True
+
+    @classmethod
+    def get_logger(cls, name: str, level: str = None) -> logging.Logger:
+        """
+        Get a logger for the given layer
+
+        Args:
+            name: Logger name (usually module or class name)
+            level: Layer tag (ENV, ROBOT, CONTROLLER, MODEL, UTILS)
+
+        Returns:
+            Configured logger instance
+        """
+        if not cls._initialized:
+            cls.setup()
+
+        # Auto-detect layer
+        if level is None:
+            level = cls._detect_level(name)
+
+        logger_key = f"{level}.{name}"
+
+        if logger_key in cls._loggers:
+            return cls._loggers[logger_key]
+
+        # Create new logger
+        logger = logging.getLogger(logger_key)
+        logger.setLevel(cls.log_level)
+        logger.propagate = False
+
+        # Clear existing handlers
+        logger.handlers.clear()
+
+        # Console output
+        if cls.console_output:
+            console_handler = logging.StreamHandler(sys.stdout)
+            console_handler.setLevel(cls.log_level)
+
+            if cls.colorful:
+                # Colored formatter
+                color = cls.LEVEL_COLORS.get(level, "white")
+                console_format = (
+                    f"%(log_color)s[%(asctime)s]%(reset)s "
+                    f"%(bold_{color})s[{level:^10}]%(reset)s "
+                    f"%(bold_white)s[%(name)s]%(reset)s "
+                    f"%(log_color)s%(levelname)-8s%(reset)s "
+                    f"%(message)s"
+                )
+
+                console_formatter = colorlog.ColoredFormatter(
+                    console_format,
+                    datefmt="%H:%M:%S",
+                    log_colors={
+                        "DEBUG": "cyan",
+                        "INFO": "green",
+                        "WARNING": "yellow",
+                        "ERROR": "red",
+                        "CRITICAL": "bold_red",
+                    },
+                )
+            else:
+                # Plain formatter
+                console_format = (
+                    f"[%(asctime)s] [{level:^10}] [%(name)s] "
+                    f"%(levelname)-8s %(message)s"
+                )
+                console_formatter = logging.Formatter(
+                    console_format, datefmt="%H:%M:%S"
+                )
+
+            console_handler.setFormatter(console_formatter)
+            logger.addHandler(console_handler)
+
+        # File output
+        if cls.file_output and cls.log_file:
+            file_handler = logging.FileHandler(cls.log_file, encoding="utf-8")
+            file_handler.setLevel(cls.log_level)
+
+            file_format = (
+                f"[%(asctime)s] [{level:^10}] [%(name)s] "
+                f"%(levelname)-8s %(message)s"
+            )
+            file_formatter = logging.Formatter(file_format, datefmt="%Y-%m-%d %H:%M:%S")
+            file_handler.setFormatter(file_formatter)
+            logger.addHandler(file_handler)
+
+        cls._loggers[logger_key] = logger
+        return logger
+
+    @classmethod
+    def _detect_level(cls, name: str) -> str:
+        """Auto-detect layer from name"""
+        name_lower = name.lower()
+
+        if "env" in name_lower:
+            return cls.LEVEL_ENV
+        elif "robot" in name_lower and "controller" not in name_lower:
+            return cls.LEVEL_ROBOT
+        elif (
+            "controller" in name_lower
+            or "communication" in name_lower
+            or "socket" in name_lower
+        ):
+            return cls.LEVEL_CONTROLLER
+        elif "model" in name_lower or "wrapper" in name_lower:
+            return cls.LEVEL_MODEL
+        else:
+            return cls.LEVEL_UTILS
+
+    @classmethod
+    def get_env_logger(cls, name: str = "Environment") -> logging.Logger:
+        """Get ENV layer logger"""
+        return cls.get_logger(name, cls.LEVEL_ENV)
+
+    @classmethod
+    def get_robot_logger(cls, name: str = "Robot") -> logging.Logger:
+        """Get ROBOT layer logger"""
+        return cls.get_logger(name, cls.LEVEL_ROBOT)
+
+    @classmethod
+    def get_controller_logger(cls, name: str = "Controller") -> logging.Logger:
+        """Get CONTROLLER layer logger"""
+        return cls.get_logger(name, cls.LEVEL_CONTROLLER)
+
+    @classmethod
+    def get_model_logger(cls, name: str = "Model") -> logging.Logger:
+        """Get MODEL layer logger"""
+        return cls.get_logger(name, cls.LEVEL_MODEL)
+
+    @classmethod
+    def get_utils_logger(cls, name: str = "Utils") -> logging.Logger:
+        """Get UTILS layer logger"""
+        return cls.get_logger(name, cls.LEVEL_UTILS)
+
+    @classmethod
+    def set_level(cls, level: str):
+        """Change log level for all loggers"""
+        new_level = getattr(logging, level.upper())
+        cls.log_level = new_level
+        for logger in cls._loggers.values():
+            logger.setLevel(new_level)
+            for handler in logger.handlers:
+                handler.setLevel(new_level)
+
+    @classmethod
+    def close_all(cls):
+        """Close file handles for all loggers"""
+        for logger in cls._loggers.values():
+            for handler in logger.handlers[:]:
+                handler.close()
+                logger.removeHandler(handler)
+        cls._loggers.clear()
+        cls._initialized = False
+
+
+# Convenience functions
+def get_logger(name: str, level: str = None) -> logging.Logger:
+    """
+    Convenience wrapper to get a logger
+
+    Args:
+        name: Logger name (usually __name__)
+        level: Layer tag (optional; auto-detected if omitted)
+
+    Returns:
+        Configured logger instance
+
+    Example:
+        from wall_x._vendor.harrix.serving._wallx_infer.logger import get_logger
+        logger = get_logger(__name__)  # auto-detect layer
+        logger = get_logger(__name__, "ROBOT")  # explicit layer
+    """
+    return InferLogger.get_logger(name, level)
+
+
+def setup_logger(
+    log_level: str = "INFO",
+    log_dir: Optional[str] = None,
+    console_output: bool = True,
+    file_output: bool = True,
+    colorful: bool = True,
+):
+    """
+    Convenience wrapper to configure logging
+
+    Args:
+        log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        log_dir: Log file directory
+        console_output: Whether to log to console
+        file_output: Whether to log to file
+        colorful: Colored console output
+
+    Example:
+        from wall_x._vendor.harrix.serving._wallx_infer.logger import setup_logger
+        setup_logger(log_level="DEBUG", log_dir="./logs")
+    """
+    InferLogger.setup(log_level, log_dir, console_output, file_output, colorful)

wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py

@@ -0,0 +1,760 @@
+import os
+import torch
+import copy
+from safetensors.torch import load_file
+import numpy as np
+from PIL import Image
+from qwen_vl_utils.vision_process import smart_resize
+from transformers import BatchFeature
+
+from wall_x.trainer.trainer_utils import load_wallx_processors
+
+from wall_x._vendor.x2robot_utils.text_templates import (
+    preprocesser_call,
+    get_prologue_with_embodied_information,
+)
+from wall_x._vendor.x2robot_utils.grounding import (
+    reverse_grounding_points,
+    extract_grounding_points,
+)
+
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from wall_x._vendor.harrix.utils.ckpt_load import reshape_compatible_state_dict
+from wall_x._vendor.harrix.utils.train_config import (
+    resolve_camera_label,
+    resolve_max_length,
+    resolve_state_bins,
+    resolve_use_state_string_representation,
+)
+from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger
+from wall_x.utils.timers import timer, ScopeTimer
+
+
+ENABLE_FAST_PREPROCESS = os.getenv("ENABLE_FAST_PREPROCESS", "False").lower() == "true"
+
+
+def move_to_cuda(obj, device="cuda"):
+    if isinstance(obj, torch.Tensor):
+        return obj.to(device)
+    elif isinstance(obj, (dict, BatchFeature)):
+        return {k: move_to_cuda(v, device) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [move_to_cuda(v, device) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(move_to_cuda(v, device) for v in obj)
+    else:
+        return obj
+
+
+class WallxModelWrapper:
+    def __init__(self, config: InferConfig):
+        self.config = config
+        self.logger = InferLogger.get_model_logger("WallxModelWrapper")
+        self.norm_key = self.config.norm_key
+        self._register_normalizers()
+        self.logger.info(f"normalizers {self.norm_key} registered")
+        self._load_processor()
+        self._load_model()
+        self.load_ckpt()
+        self.logger.info(f"model {self.config.checkpoint_path} loaded")
+
+        self.norm_key = self.config.norm_key
+        self._register_normalizers()
+        self.logger.info(f"normalizers {self.norm_key} registered")
+
+        # Initialize robot_type_id (for v3.1 delta tokenizer)
+        self.robot_type_id = None
+        if self.tokenizer_mixin is not None:
+            robot_type = getattr(self.config, "robot_type", None)
+            if robot_type:
+                self.tokenizer_mixin.init_inference(robot_type=robot_type)
+                if hasattr(self.tokenizer_mixin, "robot_type_id"):
+                    self.robot_type_id = self.tokenizer_mixin.robot_type_id
+                    if self.robot_type_id is not None:
+                        self.logger.info(
+                            f"Initialized robot_type_id: {self.robot_type_id} from robot_type: {robot_type}"
+                        )
+
+        self.role_start_symbol = "<|im_start|>"
+        self.role_end_symbol = "<|im_end|>"
+        self.vision_start_symbol = "<|vision_start|>"
+        self.vision_end_symbol = "<|vision_end|>"
+        self.image_pad_symbol = "<|image_pad|>"
+        self.propri_symbol = "<|propri|>"
+        self.action_symbol = "<|action|>"
+
+        self.cam_names = self.config.cam_names
+
+    def _camera_name_mapping(self):
+        return self.config.train_config.get("data", {}).get("camera_name_mapping")
+
+    def _load_processor(self):
+        # Load tokenizer on model_device for inference
+        device = getattr(self.config, "model_device", "cuda")
+        processors_dict = load_wallx_processors(self.config.train_config, device=device)
+        self.processor = processors_dict["processor"]
+        self.action_mapper = processors_dict["action_mapper"]
+        self.tokenizer_mixin = processors_dict.get("tokenizer_mixin")
+
+    def _load_model(self):
+        from wall_x.trainer.adapters import resolve_adapter
+
+        model_type = self.config.train_config["model_type"]
+        adapter_cls = resolve_adapter(model_type)
+        ModelClass = adapter_cls.inference_model_class()
+        self.ModelClass = ModelClass
+
+        self.logger.info(f"initializing model: {model_type} ({ModelClass.__name__})")
+
+        self.model = ModelClass(
+            self.config.model_config,
+            self.processor,
+            self.tokenizer_mixin,
+        )
+
+        # log attention implementation - variant-specific layout dispatched
+        # via the adapter.
+        adapter_cls.log_attention_implementation(self.logger, self.model)
+        self.model_type = model_type
+
+        self.logger.info("resizing model token embeddings")
+        self.model.resize_token_embeddings(len(self.processor.tokenizer))
+        self.logger.info("token embedding resize done")
+        self.logger.info("casting selected params to bfloat16")
+        self.model.to_bfloat16_for_selected_params()
+        self.logger.info("bfloat16 cast done")
+
+    def load_ckpt(self, checkpoint_path: str = None):
+        if checkpoint_path is None:
+            checkpoint_path = self.config.checkpoint_path
+
+        if os.path.exists(os.path.join(checkpoint_path, "global_step.pth")):
+            global_step = torch.load(os.path.join(checkpoint_path, "global_step.pth"))[
+                "global_step"
+            ]
+            self.logger.info(f"checkpoint global_step: {global_step}")
+
+        fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin")
+        safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors")
+        if os.path.exists(fsdp_ckpt):
+            self.logger.info(f"Loading FSDP checkpoint: {fsdp_ckpt}")
+            state_dict = torch.load(fsdp_ckpt, map_location="cpu")
+            # Unwrap outer dict (e.g. {'state_dict': {...}})
+            if isinstance(state_dict, dict) and "state_dict" in state_dict:
+                self.logger.info(
+                    "Using nested state_dict inside pytorch_model_fsdp.bin"
+                )
+                state_dict = state_dict["state_dict"]
+        elif os.path.exists(safetensor_ckpt):
+            self.logger.info(f"loading safetensors checkpoint: {safetensor_ckpt}")
+            state_dict = load_file(safetensor_ckpt, device="cpu")
+        else:
+            raise FileNotFoundError(
+                f"ERROR: No checkpoint found under {checkpoint_path}. "
+                "Expecting either pytorch_model_fsdp.bin or model.safetensors."
+            )
+
+        # Qwen models need fused weight conversion
+        if not self.ModelClass.is_fused(state_dict):
+            self.logger.info(
+                "Converting non-fused weights to fused format...",
+            )
+            state_dict = self.ModelClass.convert_to_fused(state_dict)
+        else:
+            self.logger.info(
+                "The weights is fused, skipping conversion.",
+            )
+
+        state_dict = reshape_compatible_state_dict(
+            state_dict, self.model.state_dict(), log_fn=self.logger.info
+        )
+        msg = self.model.load_state_dict(state_dict, strict=False)
+        self.model.set_normalizer(
+            copy.deepcopy(self.normalizer_action),
+            copy.deepcopy(self.normalizer_propri),
+        )
+        self.logger.info(f"load_state_dict result: {msg}")
+        self.model.eval()
+        self.model.to(self.config.model_device)
+        self.model.to_bfloat16_for_selected_params()
+
+        if hasattr(self.model, "load_optimized_weights"):
+            self.model.load_optimized_weights(state_dict)
+
+    def _register_normalizers(self):
+        from wall_x._vendor.harrix.utils.normalizer import build_normalizers
+
+        self.normalizer_action, self.normalizer_propri, resolved = build_normalizers(
+            self.config.checkpoint_path,
+            self.config.train_config,
+            self.config.norm_key,
+        )
+        self.norm_key = resolved
+        self.config.norm_key = resolved
+        self._log_norm_debug(self.norm_key)
+
+    def _log_norm_debug(self, norm_key):
+        if not hasattr(self, "normalizer_action") or not hasattr(
+            self, "normalizer_propri"
+        ):
+            self.logger.warning("[NormDebug] normalizer is not initialized yet")
+            return
+
+        if (
+            norm_key in self.normalizer_action.min
+            and norm_key in self.normalizer_propri.min
+        ):
+            self.logger.debug(
+                "[NormDebug] norm_key=%s action(min/delta) shape=%s/%s",
+                norm_key,
+                tuple(self.normalizer_action.min[norm_key].shape),
+                tuple(self.normalizer_action.delta[norm_key].shape),
+            )
+            self.logger.debug(
+                "[NormDebug] action min(all)=%s delta(all)=%s",
+                self.normalizer_action.min[norm_key].detach().cpu().tolist(),
+                self.normalizer_action.delta[norm_key].detach().cpu().tolist(),
+            )
+            self.logger.debug(
+                "[NormDebug] propri(min/delta) shape=%s/%s",
+                tuple(self.normalizer_propri.min[norm_key].shape),
+                tuple(self.normalizer_propri.delta[norm_key].shape),
+            )
+            self.logger.debug(
+                "[NormDebug] propri min(all)=%s delta(all)=%s",
+                self.normalizer_propri.min[norm_key].detach().cpu().tolist(),
+                self.normalizer_propri.delta[norm_key].detach().cpu().tolist(),
+            )
+        else:
+            self.logger.warning(
+                "[NormDebug] norm_key=%s not found in normalizer", norm_key
+            )
+
+    @timer
+    def construct_model_input(
+        self, observation, prefix_text, postfix_text, pad_prefix=False
+    ):
+        batch_size = len(observation)
+        dataset_names = [self.norm_key] * batch_size
+        self.logger.debug("[NormDebug] dataset_names=%s", dataset_names)
+
+        additional_inputs = {}
+
+        # -------- proprioception / masks (batch) --------
+        agent_pos_list = []
+        agent_pos_mask_list = []
+        dof_mask_list = []
+        for obs in observation:
+            if "robot_state_action_data" in obs:
+                robot_state_action_data = obs["robot_state_action_data"]
+
+                agent_pos = torch.from_numpy(robot_state_action_data.agent_pos)
+                # Normalize to [1, T, D] for batch cat
+                if agent_pos.dim() == 2:
+                    agent_pos = agent_pos.unsqueeze(0)
+                agent_pos_list.append(agent_pos)
+
+                agent_pos_mask = torch.from_numpy(
+                    robot_state_action_data.agent_pos_mask
+                )
+                if agent_pos_mask.dim() == 2:
+                    agent_pos_mask = agent_pos_mask.unsqueeze(0)
+                agent_pos_mask_list.append(agent_pos_mask)
+
+                dof_mask = torch.from_numpy(robot_state_action_data.dof_mask)
+                if dof_mask.dim() == 1:
+                    dof_mask = dof_mask.unsqueeze(0)
+                dof_mask_list.append(dof_mask)
+
+        # cat: [B, T, D] / [B, ...]
+        if len(agent_pos_list) > 0:
+            agent_pos = torch.cat(agent_pos_list, dim=0)
+            agent_pos_mask = torch.cat(agent_pos_mask_list, dim=0)
+            dof_mask = torch.cat(dof_mask_list, dim=0)
+
+            if self.normalizer_propri is not None:
+                agent_pos = self.normalizer_propri.normalize_data(
+                    agent_pos, dataset_names
+                )
+
+            additional_inputs["proprioception"] = agent_pos.detach()
+            additional_inputs["agent_pos_mask"] = agent_pos_mask
+            additional_inputs["dof_mask"] = dof_mask
+
+        # -------- images (flattened, in placeholder scan order) --------
+        with ScopeTimer("resize_images"):
+            # TODO[KC]: optimize this using batch processing
+            image_inputs = []
+            all_image_sizes = []
+            if ENABLE_FAST_PREPROCESS:
+                for obs in observation:
+                    current_image_inputs = self._resize_images_fast(obs)
+                    image_inputs.extend(current_image_inputs)
+                    # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format
+                    all_image_sizes.extend(
+                        [
+                            (image_i.shape[1], image_i.shape[0])
+                            for image_i in current_image_inputs
+                        ]
+                    )
+            else:
+                for obs in observation:
+                    current_image_inputs = self._resize_images(obs)
+                    image_inputs.extend(current_image_inputs)
+                    # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format
+                    all_image_sizes.extend(
+                        [
+                            (image_i.shape[1], image_i.shape[0])
+                            for image_i in current_image_inputs
+                        ]
+                    )
+
+        additional_inputs["image_size"] = all_image_sizes
+
+        with ScopeTimer("preprocesser_call"):
+            inputs = preprocesser_call(
+                processor=self.model.processor,
+                prefix_text=prefix_text,
+                postfix_text=postfix_text,
+                images=image_inputs,
+                videos=None,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+                max_length=resolve_max_length(self.config.train_config),
+                pad_to_128_multiple=False,
+                pad_prefix_to_same_length=pad_prefix,
+                norm_state=(
+                    additional_inputs["proprioception"]
+                    if "proprioception" in additional_inputs
+                    and resolve_use_state_string_representation(
+                        self.config.train_config
+                    )
+                    else None
+                ),
+                agent_pos_mask=(
+                    additional_inputs["agent_pos_mask"]
+                    if "agent_pos_mask" in additional_inputs
+                    else None
+                ),
+                state_augmentation_prob=0.0,
+                state_drop_prob=0.0,
+                state_augmentation_ratio=0.0,
+                state_bins=resolve_state_bins(self.config.train_config),
+                inference_mode=True,
+            )
+
+        with ScopeTimer("convert_action_token_id and post "):
+            action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids(
+                "<|action|>"
+            )
+            flow_action_mask = inputs["input_ids"] == action_token_id
+            additional_inputs["moe_token_types"] = flow_action_mask
+            additional_inputs["dataset_names"] = dataset_names
+
+            inputs.update(additional_inputs)
+            inputs = move_to_cuda(inputs, self.config.model_device)
+
+        return inputs
+
+    def get_text_for_dllm_action(self, instruction):
+        if (
+            self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0)
+            > 0
+        ):
+            if self.norm_key != "x2_normal" and self.norm_key != "ex_normal":
+                self.config.robot_id = 0
+            cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names}
+            prologue = get_prologue_with_embodied_information(
+                dataset_name=self.norm_key,
+                cam_mapping=cam_name_mapping,
+                robot_id=self.config.robot_id,
+                uid="",
+                config=self.config.data_config,
+            )
+        else:
+            prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        placeholder_seq = self.tokenizer_mixin.get_placeholder_for_dllm()
+        ar_token = "".join(placeholder_seq) + "<|im_end|>\n"
+        assistant_message = f"{self.role_start_symbol}assistant\n{ar_token}"
+        flow_action = f"{self.action_symbol * self.config.action_horizon}"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = flow_action
+
+        return prefix_text, postfix_text
+
+    @timer
+    def get_text_for_action(self, instruction):
+        if (
+            self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0)
+            > 0
+        ):
+            if self.norm_key not in ["x2_normal", "ex_normal"]:
+                self.config.robot_id = 0
+            cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names}
+            prologue = get_prologue_with_embodied_information(
+                dataset_name=self.norm_key,
+                cam_mapping=cam_name_mapping,
+                robot_id=self.config.robot_id,
+                uid="",
+                config=self.config.data_config,
+            )
+        else:
+            prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        assistant_message = f"{self.role_start_symbol}assistant\n"
+        flow_action = f"{self.action_symbol * self.config.action_horizon}"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = flow_action
+
+        return prefix_text, postfix_text
+
+    @timer
+    def get_text_for_subtask_generation(self, instruction):
+        prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = "\nPredict the next action in language.\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        assistant_message = f"{self.role_start_symbol}assistant\n"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = ""
+
+        return prefix_text, postfix_text
+
+    @timer
+    def _resize_images(self, observation):
+        image_inputs = []
+        for key in self.cam_names:
+            if key not in observation:
+                continue
+            current_obs = observation[key]
+            if isinstance(current_obs, np.ndarray):
+                img_pil = Image.fromarray(current_obs)
+            elif isinstance(current_obs, Image.Image):
+                img_pil = current_obs
+            else:
+                raise ValueError(f"Unsupported image type: {type(current_obs)}")
+            orig_width, orig_height = img_pil.size
+
+            target_size = self.config.data_config.resolution.get(key, -1)
+            if target_size != -1:
+                # Aspect-ratio-preserving resize
+                if orig_width > orig_height:  # landscape
+                    new_width = target_size
+                    new_height = int(target_size * orig_height / orig_width)
+                else:  # portrait
+                    new_height = target_size
+                    new_width = int(target_size * orig_width / orig_height)
+                img_pil = img_pil.resize((new_width, new_height))
+
+            # Apply smart resize (Qwen logic)
+            current_width, current_height = img_pil.size
+            resized_height, resized_width = smart_resize(
+                current_height,
+                current_width,
+                factor=self.config.data_config.image_factor,
+                min_pixels=self.config.data_config.min_pixels,
+                max_pixels=self.config.data_config.max_pixels,
+            )
+            resized_img = img_pil.resize((resized_width, resized_height))
+            resized_img = torch.from_numpy(np.array(resized_img)).to(
+                self.config.model_device
+            )
+            image_inputs.append(resized_img)
+
+        return image_inputs
+
+    def _resize_images_fast(self, observation):
+        import cv2
+
+        image_inputs = []
+        for key in self.cam_names:
+            if key not in observation:
+                continue
+            current_obs = observation[key]
+            orig_height, orig_width, _ = current_obs.shape
+
+            target_size = self.config.data_config.resolution.get(key, -1)
+            current_width, current_height = orig_width, orig_height
+            if target_size != -1:
+                # Aspect-ratio-preserving resize
+                if orig_width > orig_height:  # landscape
+                    new_width = target_size
+                    new_height = int(target_size * orig_height / orig_width)
+                else:  # portrait
+                    new_height = target_size
+                    new_width = int(target_size * orig_width / orig_height)
+                current_width = new_width
+                current_height = new_height
+
+            # Apply smart resize (Qwen logic)
+            resized_height, resized_width = smart_resize(
+                current_height,
+                current_width,
+                factor=self.config.data_config.image_factor,  # FIXME
+                min_pixels=self.config.data_config.min_pixels,  # FIXME
+                max_pixels=self.config.data_config.max_pixels,  # FIXME
+            )
+
+            resized_img = cv2.resize(
+                current_obs,
+                (resized_width, resized_height),
+                interpolation=cv2.INTER_CUBIC,
+            )
+            resized_img = torch.from_numpy(resized_img).to(self.config.model_device)
+            image_inputs.append(resized_img)
+
+        return image_inputs
+
+    def infer_flow_action(self, observation, instruction):
+        self.logger.info("generating flow action")
+        self.logger.info(f"flow action instruction: {instruction}")
+
+        prefix_text, postfix_text = self.get_text_for_action(instruction)
+        model_input = self.construct_model_input(
+            [observation], [prefix_text], [postfix_text]
+        )
+
+        padding = (
+            torch.zeros_like(
+                self.normalizer_action.delta[model_input["dataset_names"][0]]
+            )
+            .unsqueeze(0)
+            .to("cpu")
+        )
+        padding_action = self.normalizer_action.normalize_data(
+            padding, model_input["dataset_names"]
+        ).to(model_input["input_ids"].device)
+
+        self.logger.info(
+            "generate_flow_action start (horizon=%s, flow_steps=%s, device=%s)",
+            self.config.action_horizon,
+            self.config.num_inference_timesteps,
+            self.config.model_device,
+        )
+        if torch.cuda.is_available():
+            try:
+                free_b, total_b = torch.cuda.mem_get_info(
+                    torch.device(self.config.model_device)
+                )
+                self.logger.info(
+                    "CUDA mem before flow: free=%.2f GiB / total=%.2f GiB",
+                    free_b / (1024**3),
+                    total_b / (1024**3),
+                )
+            except Exception as e:
+                self.logger.warning("CUDA mem_get_info failed: %s", e)
+
+        with ScopeTimer("generate_flow_action"):
+            model_output = self.model.generate_flow_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                padding_action=padding_action,
+                **model_input,
+            )
+
+        self.logger.info("flow action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved flow action to robot_state_action_data")
+        return model_output
+
+    def infer_flow_action_batch(self, observations, instructions):
+        """
+        Batch flow action inference:
+        - observations: List[Dict], same format as single inference
+        - instructions: List[str], aligned with observations
+        Returns List[model_output] of length batch size
+        """
+        assert len(observations) == len(
+            instructions
+        ), "observations and instructions must have the same length"
+        batch_size = len(observations)
+
+        prefix_list = []
+        postfix_list = []
+        for ins in instructions:
+            prefix_text, postfix_text = self.get_text_for_action(ins)
+            prefix_list.append(prefix_text)
+            postfix_list.append(postfix_text)
+
+        # Build batch input in one preprocesser_call; avoid per-sample cat
+        batch_inputs = self.construct_model_input(
+            observations, prefix_list, postfix_list
+        )
+
+        padding_list = []
+        for ds_name in batch_inputs["dataset_names"]:
+            padding = (
+                torch.zeros_like(self.normalizer_action.delta[ds_name])
+                .unsqueeze(0)
+                .to("cpu")
+            )
+            padding_list.append(padding)
+        padding = torch.cat(padding_list, dim=0)
+        padding_action = self.normalizer_action.normalize_data(
+            padding, batch_inputs["dataset_names"]
+        ).to(batch_inputs["input_ids"].device)
+
+        with ScopeTimer("generate_flow_action_batch"):
+            model_output = self.model.generate_flow_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                padding_action=padding_action,
+                **batch_inputs,
+            )
+
+        predict_action = model_output["predict_action"]  # [B, H, D]
+
+        outputs = []
+        for i in range(batch_size):
+            single_action = predict_action[i : i + 1]
+            single_output = {
+                "predict_action": single_action,
+                "robot_state_action_data": observations[i]["robot_state_action_data"],
+            }
+            single_output["robot_state_action_data"].save_action_data(
+                single_output["predict_action"]
+            )
+            outputs.append(single_output)
+
+        return outputs
+
+    def infer_ar_action(self, observation, instruction):
+        self.logger.info("generating ar action")
+        self.logger.info(f"ar action instruction: {instruction}")
+
+        prefix_text, _ = self.get_text_for_action(instruction)
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_ar_action(
+            action_horizon=self.config.action_horizon,
+            action_dim=self.config.ar_action_dim,
+            num_inference_timesteps=self.config.num_inference_timesteps,
+            robot_type_id=self.robot_type_id,
+            **model_input,
+        )
+        self.logger.info("ar action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved ar action to robot_state_action_data")
+        return model_output
+
+    def infer_subtask(self, observation, instruction):
+        self.logger.info("generating subtask")
+        self.logger.info(f"subtask instruction: {instruction}")
+
+        prefix_text, postfix_text = self.get_text_for_subtask_generation(instruction)
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_text(**model_input)
+        subtask = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip()
+        self.logger.info(f"subtask generation done, subtask: {subtask}")
+        return subtask
+
+    def infer_vqa(self, observation, instruction):
+        self.logger.info("generating vqa answer")
+
+        if isinstance(observation["multi_modal"], list):
+            orig_size = observation["multi_modal"][0].size
+        else:
+            orig_size = observation["multi_modal"].size
+
+        prefix_text = instruction
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_text(**model_input)
+        answer = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip()
+        self.logger.info(f"vqa answer done, answer: {answer}")
+        answer = reverse_grounding_points(
+            answer,
+            orig_size[1],
+            orig_size[0],
+            model_input["image_size"][0][1],
+            model_input["image_size"][0][0],
+            self.config.data_config.model_type,
+        )
+        points = extract_grounding_points(answer)
+
+        return {
+            "answer": answer,
+            "points": points,
+        }
+
+    def infer_dllm_action(
+        self, observation, instruction, use_ar_action=False, dataset_name="x2_normal"
+    ):
+        assert self.model_type in ["qwen2_5"], "DLLM only supports qwen2_5"
+        prefix_text, postfix_text = self.get_text_for_dllm_action(instruction)
+        model_input = self.construct_model_input(
+            [observation], [prefix_text], [postfix_text], pad_prefix=True
+        )
+        model_input = self.model.update_infer_dllm_position_mask(model_input)
+        total_ar_step = self.tokenizer_mixin.inference_ar_steps_for_dllm
+        cnt = 0
+        while cnt < 3:  # fast mode: at most 3 retries
+            model_output = self.model.generate_dllm_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                ar_action_dim=self.config.ar_action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                use_ar_action=use_ar_action,
+                total_ar_step=total_ar_step,
+                robot_type_id=self.robot_type_id,  # for v3.1 delta decode
+                **model_input,
+            )
+            if model_output["predict_action"] is not None:
+                break
+            cnt += 1
+            self.logger.warning(f"dllm action generation failed, retry {cnt}")
+        self.logger.info("dllm action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved dllm action to robot_state_action_data")
+        return model_output

wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py

@@ -0,0 +1,423 @@
+import socket
+import struct
+import json
+import cv2
+import numpy as np
+import time
+import select
+import errno
+import os
+import pickle
+
+import threading
+
+from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger
+
+_DEFAULT_ROBOT_CONFIG = {
+    0: {"host": "127.0.0.1", "action_port": 57749, "keyboard_port": 58849},
+    1: {"host": "127.0.0.1", "action_port": 57750, "keyboard_port": 58850},
+    2: {"host": "127.0.0.1", "action_port": 57751, "keyboard_port": 58851},
+    3: {"host": "127.0.0.1", "action_port": 57761, "keyboard_port": 58861},
+}
+
+
+class RobotRegistry:
+    def __init__(self):
+        self.registry = _DEFAULT_ROBOT_CONFIG
+        self.logger = InferLogger.get_controller_logger("RobotRegistry")
+        self.logger.debug(f"init robot registry, registered {len(self.registry)} robots")
+
+    def register_robot(self, robot_id, host, action_port, keyboard_port):
+        if robot_id in self.registry:
+            self.logger.warning(f"robot ID {robot_id} already registered")
+            return False
+        else:
+            self.registry[robot_id] = {
+                "host": host,
+                "action_port": action_port,
+                "keyboard_port": keyboard_port,
+            }
+            self.logger.info(
+                f"register robot ID={robot_id}: host={host}, action_port={action_port}, keyboard_port={keyboard_port}"
+            )
+            return True
+
+    def get_robot_info(self, robot_id):
+        if robot_id in self.registry:
+            info = self.registry[robot_id]
+            self.logger.debug(f"get robot ID={robot_id} info: {info}")
+            return info
+        else:
+            self.logger.error(f"robot ID={robot_id} not found")
+            return None
+
+    def exist(self, robot_id):
+        return robot_id in self.registry
+
+
+class RobotCommunication:
+    def __init__(self, robot_id, host=None, action_port=None, keyboard_port=None):
+        self.logger = InferLogger.get_controller_logger("RobotCommunication")
+        self.logger.debug(f"init robot comms: robot_id={robot_id}")
+
+        self.robot_register = RobotRegistry()
+        if self.robot_register.exist(robot_id):
+            self.robot_info = self.robot_register.get_robot_info(robot_id)
+            self.logger.debug(f"using registered robot config: {robot_id}")
+        else:
+            self.robot_register.register_robot(
+                robot_id, host, action_port, keyboard_port
+            )
+            self.robot_info = self.robot_register.get_robot_info(robot_id)
+            self.logger.debug(f"register new robot config: {robot_id}")
+
+        self.action_sock = None
+        self.action_conn = None
+
+        self.keyboard_sock = None
+        self.keyboard_conn = None
+
+        self.client_socks = []
+
+    def connect(self):
+        self.logger.info("establishing socket connection...")
+
+        self.action_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.action_sock.setblocking(True)
+        self.action_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+
+        host = self.robot_info["host"]
+        port = self.robot_info["action_port"]
+        self.action_sock.bind((host, port))
+        self.action_sock.listen(1)
+        self.logger.info(f"listening on action port: {host}:{self.action_sock.getsockname()[1]}")
+
+        action_thread = threading.Thread(target=self.handle_action_client)
+
+        self.logger.debug("starting connection handler thread...")
+        action_thread.start()
+
+        action_thread.join()
+        self.logger.info("socket connection established")
+
+    def handle_action_client(self):
+        self.logger.debug("waiting for client connection...")
+        self.action_conn, addr = self.action_sock.accept()
+        self.logger.info(f"accepted connection from {addr}")
+
+    def recv_image(self, index):
+        self.logger.debug(f"receiving image index={index}...")
+        image_size = struct.unpack("<L", self.action_conn.recv(4))[0]
+        self.logger.debug(f"image size: {image_size} bytes")
+
+        image = self.recvall(self.action_conn, image_size)
+        nparr = np.frombuffer(image, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        self.logger.debug(f"image decoded: shape={image.shape}")
+        return image
+
+    def recv_action_data(self):
+        self.logger.debug("receiving action data...")
+        data_size = struct.unpack("<L", self.action_conn.recv(4))[0]
+        self.logger.debug(f"action data size: {data_size} bytes")
+
+        data = self.recvall(self.action_conn, data_size)
+        action_data = json.loads(data.decode("utf8"))
+
+        self.logger.debug(f"action data received: {len(action_data)} fields")
+        return action_data
+
+    def accept_connections(self):
+        try:
+            client_sock, addr = self.keyboard_sock.accept()
+            client_sock.setblocking(0)  # Set non-blocking mode for the client socket
+            self.client_socks.append(client_sock)
+        except BlockingIOError:
+            # The socket is in non-blocking mode and there are no pending connections
+            pass
+
+    def recv_keyboard_input(self):
+        for client_sock in list(self.client_socks):  # Iterate over a copy of the list
+            if client_sock.fileno() == -1:
+                # Socket has been closed, remove it from the list
+                self.client_socks.remove(client_sock)
+                continue
+            read_sockets, _, _ = select.select([client_sock], [], [], 0)
+            if client_sock in read_sockets:
+                try:
+                    data_size_bytes = client_sock.recv(4)
+                    if not data_size_bytes:
+                        self.logger.debug("client closed connection")
+                        client_sock.close()
+                        self.client_socks.remove(client_sock)
+                        continue
+                    data_size = struct.unpack("<L", data_size_bytes)[0]
+                    data = RobotCommunication.recvall(client_sock, data_size)
+                    if not data:
+                        self.logger.debug("client closed connection")
+                        client_sock.close()
+                        self.client_socks.remove(client_sock)
+                        continue
+                    json_data = json.loads(data.decode("utf8"))
+                    self.logger.debug(f"keyboard input: {json_data}")
+                    return json_data
+                except socket.error as e:
+                    if e.errno == errno.EAGAIN or e.errno == errno.EWOULDBLOCK:
+                        # Resource temporarily unavailable, ignore the error
+                        pass
+                    else:
+                        self.logger.error(f"socket error: {e}")
+                        client_sock.close()
+                        self.client_socks.remove(client_sock)
+                        continue
+        time.sleep(0.1)
+        return None
+
+    def send_dict(self, dict_data):
+        self.logger.debug(f"sending dict: {len(dict_data)} fields")
+        data_str = json.dumps(dict_data)
+        data_bytes = data_str.encode("utf-8")
+        self.action_conn.sendall(struct.pack("<L", len(data_bytes)))
+        self.action_conn.sendall(data_bytes)
+        self.logger.debug(f"data sent: {len(data_bytes)} bytes")
+
+    def close(self):
+        self.logger.info("closing socket connection...")
+        self.action_sock.close()
+        for sock in self.client_socks:
+            sock.close()
+        self.logger.info(f"closed {len(self.client_socks)} client connections")
+        # self.keyboard_sock.close()
+
+    @staticmethod
+    def recvall(sock, count):
+        buf = b""
+        while count:
+            newbuf = sock.recv(count)
+            if not newbuf:
+                return None
+            buf += newbuf
+            count -= len(newbuf)
+        return buf
+
+
+class DummyRobotCommunication:
+    def __init__(self):
+        pass
+
+    def send_dict(self, dict_data):
+        pass
+
+
+class DummyRobotController:
+    def __init__(
+        self,
+        robot_id: int,
+        host: str = None,
+        port: str = None,
+        views_path: str | None = None,
+        state_path: str | None = None,
+    ):
+        self.robot_id = robot_id
+        self.host = host
+        self.port = port
+        self.views_path = views_path or os.environ.get("WALL_X_DUMMY_VIEWS_PATH")
+        self.state_path = state_path or os.environ.get("WALL_X_DUMMY_STATE_PATH")
+        self.robot_comm = DummyRobotCommunication()
+        self.logger = InferLogger.get_controller_logger("DummyRobotController")
+        self.logger.info(
+            f"init virtual robot controller (debug): robot_id={robot_id}"
+        )
+
+    def connect(self):
+        self.logger.debug("virtual controller connect (no-op)")
+        pass
+
+    def close(self):
+        self.logger.debug("virtual controller close (no-op)")
+        pass
+
+    def recv_image(
+        self, cam_names: list = ["camera_left", "camera_front", "camera_right"]
+    ) -> dict:
+        if not self.views_path:
+            raise RuntimeError(
+                "DummyRobotController requires views_path or "
+                "WALL_X_DUMMY_VIEWS_PATH."
+            )
+        self.logger.debug(f"load views from pickle: {cam_names}")
+        with open(self.views_path, "rb") as f:
+            all_views = pickle.load(f)
+        self.logger.debug(f"views loaded: {len(all_views)} views")
+        return all_views
+
+    def recv_action(self):
+        if not self.state_path:
+            raise RuntimeError(
+                "DummyRobotController requires state_path or "
+                "WALL_X_DUMMY_STATE_PATH."
+            )
+        self.logger.debug("load actions from pickle")
+        with open(self.state_path, "rb") as f:
+            all_actions = pickle.load(f)
+        self.logger.debug(f"actions loaded: {len(all_actions)} fields")
+        return all_actions
+
+
+class RobotController:
+    def __init__(
+        self,
+        robot_id: int,
+        host: str = None,
+        port: str = None,
+        max_time_step: int = 10000,
+    ):
+        self.logger = InferLogger.get_controller_logger("RobotController")
+        self.logger.info(
+            f"init robot controller: robot_id={robot_id}, max_time_step={max_time_step}"
+        )
+
+        self.robot_comm = RobotCommunication(robot_id, host, port)
+        self.max_time_step = max_time_step
+        self.global_step = 0
+
+    def connect(self):
+        self.logger.info("connecting robot controller...")
+        self.robot_comm.connect()
+        self.logger.info("robot controller connected")
+
+    def close(self):
+        self.logger.info("closing robot controller...")
+        self.robot_comm.close()
+        self.logger.info("robot controller closed")
+
+    def prediction(self, views, actions) -> dict:
+        raise NotImplementedError
+
+    def recv_image(
+        self, cam_names: list = ["camera_left", "camera_front", "camera_right"]
+    ) -> dict:
+        self.logger.debug(f"receive images: {cam_names}")
+        views = {}
+        for i, name in enumerate(cam_names):
+            image = np.array(self.robot_comm.recv_image(i))
+            views[name] = image[None, :]
+            self.logger.debug(f"received {name}: shape={image.shape}")
+        self.logger.debug(f"all images received: {len(views)} views")
+        return views
+
+    def recv_action(self):
+        self.logger.debug("receiving action data...")
+        action_data = self.robot_comm.recv_action_data()
+        return action_data
+
+    def recv_keyboard_input(self):
+        self.robot_comm.accept_connections()
+        json_data = self.robot_comm.recv_keyboard_input()
+        if json_data is not None:
+            self.logger.debug(
+                f"keyboard: motionStatus={json_data.get('motionStatus')}, armMode={json_data.get('armMode')}"
+            )
+            return json_data["motionStatus"], json_data["armMode"]
+
+        return None, None
+
+    def reset(self):
+        self.logger.info(f"reset controller: global_step {self.global_step} -> 0")
+        self.global_step = 0
+
+    def record_start(self):
+        self.logger.info("start recording...")
+        record_signal = {"cmd": "RECORD_START"}
+        self.robot_comm.send_dict(record_signal)
+
+    def record_continue(self):
+        self.logger.debug("continue recording...")
+        record_signal = {"cmd": "RECORD_CONTINUE"}
+        self.robot_comm.send_dict(record_signal)
+
+    def record_start_process(self):
+        if self.global_step == 0:
+            self.logger.info("waiting for START signal...")
+            action = None
+            while action != "START":
+                action, _ = self.recv_keyboard_input()
+            self.logger.info("received START signal")
+            self.record_start()
+
+        self.record_continue()
+
+    def set_zero(self):
+        self.logger.info("setting zero pose...")
+        record_signal = {"cmd": "INIT_ZERO", "gripper": [0.0, 0.0]}
+        self.robot_comm.send_dict(record_signal)
+        self.reset()
+        self.logger.info("zero pose set")
+
+    def record_stop(self):
+        self.logger.info("stop recording...")
+        record_signal = {"cmd": "RECORD_STOP"}
+        self.robot_comm.send_dict(record_signal)
+        time.sleep(0.01)
+
+    def recover_from_failure(self):
+        self.logger.warning("recovering from failure...")
+        record_signal = {"cmd": "TO_MASTER_SLAVE"}
+        self.robot_comm.send_dict(record_signal)
+        time.sleep(0.01)
+
+        self.logger.info("waiting for START signal...")
+        action = None
+        while action != "START":
+            action, _ = self.recv_keyboard_input()
+
+        self.logger.info("recovery done, resume recording")
+        self.record_start()
+
+    def reset_to_zero(self):
+        self.logger.info("reset to zero pose...")
+        self.record_stop()
+        self.set_zero()
+
+    def to_slave(self):
+        self.logger.info("switching to slave mode...")
+        record_signal = {"cmd": "TO_SLAVE"}
+        self.robot_comm.send_dict(record_signal)
+        time.sleep(0.01)
+
+    def run(
+        self,
+        record_mode=False,
+        cam_names: list = ["camera_left", "camera_front", "camera_right"],
+    ):
+
+        while self.global_step <= self.max_time_step:
+            if record_mode:
+                self.record_start_process()
+
+            self.global_step += 1
+
+            action = self.recv_action()
+            view = self.recv_image(cam_names)
+
+            pred = self.prediction(view, action)
+            self.robot_comm.send_dict(pred)
+
+            if record_mode:
+                action, arm_mode = self.recv_keyboard_input()
+
+                if action is not None and arm_mode is not None:
+                    self.logger.info("action: %s, arm_mode: %s", action, arm_mode)
+                    if action == "STOP" and arm_mode == "ARM_TEST_MODE_MS":
+                        self.record_stop()
+                        self.recover_from_failure()
+                        action, arm_mode = None, None
+                        while action != "STOP" or arm_mode != "ARM_TEST_MODE_S":
+                            action, arm_mode = self.recv_keyboard_input()
+                        self.record_stop()
+                        self.to_slave()
+                        self.reset()
+                    elif arm_mode == "ARM_TEST_MODE_S" and action == "INIT":
+                        self.reset_to_zero()

wall_x/_vendor/harrix/serving/_wallx_infer/utils.py

@@ -0,0 +1,280 @@
+import numpy as np
+from scipy.signal import savgol_filter
+from scipy.spatial.transform import Rotation as R  # TODO: convert to numba
+from collections import deque
+import threading
+
+from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger
+
+
+class KeyboardThread(threading.Thread):
+    """
+    Simple keyboard listener thread with stop and reset
+    """
+
+    def __init__(self):
+        self.should_reset = False
+        self.should_stop = False
+        self.new_instruction_index = None  # Stores new instruction index
+        self.logger = InferLogger.get_utils_logger("KeyboardThread")
+
+        super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True)
+        self.show_help()
+        self.start()
+
+    def run(self):
+        """Listen for keyboard input"""
+        while True:
+            try:
+                user_input = input().strip().lower()
+
+                if user_input in ["s", "stop"]:
+                    self.should_stop = not self.should_stop
+                    self.logger.info("[keyboard] stop signal sent")
+
+                elif user_input in ["r", "reset"]:
+                    self.logger.info("[keyboard] resetting...")
+                    self.should_reset = True
+                    self.logger.info("[keyboard] reset signal sent")
+
+                elif user_input.isdigit():
+                    # Numeric input: switch instruction index
+                    index = int(user_input)
+                    self.new_instruction_index = index
+                    self.logger.info(f"[keyboard] switched to instruction index: {index}")
+
+                else:
+                    self.logger.info(f"[keyboard] input received: {user_input}. no action taken.")
+
+            except EOFError:
+                break
+            except Exception as e:
+                self.logger.error(f"[keyboard] error: {e}")
+
+    def show_help(self):
+        self.logger.info(
+            "[keyboard] controls: 's' stop, 'r' reset, digit switches instruction index"
+        )
+
+
+# Arm trajectory parameters
+ARM_MAX_VELOCITY = 0.02
+ARM_EXECUTION_HZ = 20
+ARM_MIN_EXECUTION_TIME = 5.0
+ARM_MAX_EXECUTION_TIME = 15.0
+
+
+class UnifiedTrajectoryProcessor:
+    """Unified trajectory processor"""
+
+    @staticmethod
+    def interpolate_trajectory_batch(trajectories, target_length, smooth=True):
+        """
+        Interpolate multiple trajectories to a common length
+        Args:
+            trajectories: list of np.array, each shape (N, D)
+            target_length: int, target length
+            smooth: bool, whether to smooth
+        Returns:
+            list of np.array, interpolated trajectories
+        """
+        if not trajectories:
+            return []
+
+        results = []
+        for traj in trajectories:
+            if len(traj) == 0:
+                results.append(np.zeros((target_length, traj.shape[1])))
+                continue
+
+            if len(traj) == target_length:
+                results.append(traj)
+                continue
+
+            # Vectorized interpolation
+            original_indices = np.linspace(0, len(traj) - 1, len(traj))
+            target_indices = np.linspace(0, len(traj) - 1, target_length)
+
+            # Handle different data types
+            if traj.shape[1] == 7:  # Arm data [x,y,z,rx,ry,rz,gripper]
+                interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory(
+                    traj, original_indices, target_indices, target_length
+                )
+            else:  # Other data (height, current, etc.)
+                interpolated = np.zeros((target_length, traj.shape[1]))
+                for i in range(traj.shape[1]):
+                    interpolated[:, i] = np.interp(
+                        target_indices, original_indices, traj[:, i]
+                    )
+
+            # Smoothing
+            if smooth and len(interpolated) >= 5:
+                interpolated = UnifiedTrajectoryProcessor._smooth_trajectory(
+                    interpolated
+                )
+
+            results.append(interpolated)
+
+        return results
+
+    @staticmethod
+    def _interpolate_arm_trajectory(
+        traj, original_indices, target_indices, target_length
+    ):
+        """Optimized arm trajectory interpolation"""
+        interpolated = np.zeros((target_length, 7))
+
+        # Vectorized interp for position and gripper
+        for i in [0, 1, 2, 6]:  # x, y, z, gripper
+            interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
+
+        # Quaternion interpolation (vectorized)
+        quaternions = R.from_euler("xyz", traj[:, 3:6]).as_quat()
+        interpolated_quats = np.zeros((target_length, 4))
+        for i in range(4):
+            interpolated_quats[:, i] = np.interp(
+                target_indices, original_indices, quaternions[:, i]
+            )
+
+        # Batch normalize
+        norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True)
+        interpolated_quats = interpolated_quats / norms
+
+        # Batch convert back to euler
+        interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz")
+
+        return interpolated
+
+    @staticmethod
+    def _interpolate_position_trajectory(
+        traj, original_indices, target_indices, target_length
+    ):
+        """Optimized position trajectory interpolation"""
+        interpolated = np.zeros((target_length, 3))
+        for i in range(3):
+            interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
+        return interpolated
+
+    @staticmethod
+    def _smooth_trajectory(trajectory):
+        """Vectorized smoothing"""
+        if len(trajectory) < 5:
+            return trajectory
+
+        try:
+            # Smooth all dimensions in batch
+            smoothed = np.zeros_like(trajectory)
+            for dim in range(trajectory.shape[1]):
+                smoothed[:, dim] = savgol_filter(
+                    trajectory[:, dim],
+                    min(
+                        5,
+                        (
+                            len(trajectory)
+                            if len(trajectory) % 2 == 1
+                            else len(trajectory) - 1
+                        ),
+                    ),
+                    3,
+                    mode="nearest",
+                )
+            return smoothed
+        except Exception:
+            return trajectory
+
+    @staticmethod
+    def calculate_optimal_trajectory_length(left_traj, right_traj):
+        """Compute optimal trajectory length"""
+
+        # Vectorized distance computation
+        def calc_distance(traj):
+            if len(traj) < 2:
+                return 0.0
+            pos_diff = traj[1:, :3] - traj[:-1, :3]
+            return np.sum(np.linalg.norm(pos_diff, axis=1))
+
+        distances = [calc_distance(left_traj), calc_distance(right_traj)]
+        max_distance = max(distances)
+
+        if max_distance > 1e-6:
+            execution_time = np.clip(
+                max_distance / ARM_MAX_VELOCITY,
+                ARM_MIN_EXECUTION_TIME,
+                ARM_MAX_EXECUTION_TIME,
+            )
+        else:
+            execution_time = ARM_MIN_EXECUTION_TIME
+
+        return max(int(execution_time * ARM_EXECUTION_HZ), len(left_traj))
+
+
+class VehiclePoseHandler:
+    """Vehicle pose and velocity computation"""
+
+    def __init__(self):
+        self.current_pose = None
+        self.previous_pose = None
+        self.pose_history = deque(maxlen=10)
+        self.logger = InferLogger.get_utils_logger("VehiclePoseHandler")
+
+    def update_pose(self, new_pose):
+        """Update vehicle pose"""
+        if new_pose is not None:
+            self.previous_pose = self.current_pose
+            self.current_pose = np.array(new_pose)
+            self.pose_history.append(self.current_pose.copy())
+            self.logger.info("current_pose %s", self.current_pose)
+        return self.current_pose
+
+    def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None):
+        """Convert body-frame velocity to global pose"""
+        if start_pose is None:
+            if self.current_pose is not None:
+                start_pose = self.current_pose.copy()
+            else:
+                start_pose = np.array([0.0, 0.0, 0.0])
+
+        x, y, theta = start_pose
+
+        # Body velocity to global displacement
+        cos_theta = np.cos(theta)
+        sin_theta = np.sin(theta)
+
+        # Transform: body frame -> global frame
+        dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt
+        dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt
+        dtheta = vyaw * dt
+
+        # Compute new pose
+        x_new = x + dx_global
+        y_new = y + dy_global
+        theta_new = theta + dtheta
+
+        # Wrap angle to [-pi, pi]
+        theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi
+
+        return np.array([x_new, y_new, theta_new])
+
+    def compute_body_velocities_from_poses(
+        self, current_pose, previous_pose, dt=1 / 20
+    ):
+        """Compute body-frame velocity from pose delta"""
+        if current_pose is None or previous_pose is None:
+            return np.array([0.0, 0.0, 0.0])
+
+        # Global-frame displacement
+        dx_global = current_pose[0] - previous_pose[0]
+        dy_global = current_pose[1] - previous_pose[1]
+        dtheta = current_pose[2] - previous_pose[2]
+
+        # Use previous frame angle for transform
+        theta = previous_pose[2]
+        cos_theta = np.cos(theta)
+        sin_theta = np.sin(theta)
+
+        # Global displacement to body velocity
+        vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt
+        vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt
+        vyaw = dtheta / dt
+
+        return np.array([vx_body, vy_body, vyaw])

wall_x/_vendor/harrix/serving/launch_serving.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+"""
+Server script for Wall-X model.
+
+This script serves a Wall-X model using a websocket server, allowing
+clients to connect and get action predictions from observations.
+
+"""
+
+import dataclasses
+import enum
+import inspect
+import logging
+import os
+import socket
+import sys
+import yaml
+import traceback
+import tyro
+
+from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer
+
+
+def _server_model_config_to_infer_kwargs(model_config) -> dict:
+    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+
+    infer_params = inspect.signature(InferConfig.__init__).parameters
+    return {
+        k: v
+        for k, v in vars(model_config).items()
+        if v is not None and k in infer_params
+    }
+
+
+def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True):
+    from wall_x._vendor.harrix.serving.policy.wall_x_policy import WallXPolicy
+    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+
+    config = InferConfig(**_server_model_config_to_infer_kwargs(model_config))
+    return WallXPolicy(config=config, image_passing_mode=image_passing_mode, serialize_actions=serialize_actions)
+
+
+
+
+logger = logging.getLogger(__name__)
+
+
+class EnvMode(enum.Enum):
+    """Supported environments/datasets."""
+
+    X2ROBOT = "x2robot"
+    LIBERO = "libero"
+
+
+@dataclasses.dataclass
+class ServerModelConfig:
+    """Configuration for loading a Wall-X model."""
+
+    checkpoint_path: str | None = None
+    train_config_path: str | None = None
+    # robot_host: str = '0.0.0.0'
+    # robot_port: int = 33723
+    robot_type: str = "desktop"  # ["desktop", "turtle"]
+    robot_action_start_ratio: float = (
+        0.0  # proportion of action execution to start from
+    )
+    robot_action_end_ratio: float = 1.0  # proportion of action execution to end at
+    robot_action_interpolate_multiplier: int = 10  # action interpolation multiplier
+    robot_use_joint_angle_control: bool = (
+        False  # use joint angle control (model must predict joints)
+    )
+    turtle_as_desktop: bool = (
+        False  # use turtle platform as desktop with fixed base/head/camera/height
+    )
+    action_horizon: int = 32  # specify the correct horizon for the model
+    action_dim: int | None = None
+    model_device: str = "cuda"
+    num_inference_timesteps: int = 10
+    num_inference_steps: int | None = None
+    cfg_scale: float | None = None
+    seed: int | None = None
+    save_video_dir: str = "./videos"
+
+    # Please specify explicitly if the checkpoint was not trained on the x2robot dataset.
+    norm_key: str | None = None
+    # Model cameras; None = infer from train config ``data.key_mappings.camera``.
+    cam_names: list[str] | None = None
+    # Robot camera keys in incoming websocket observations.
+    camera_front_key: str = "camera_front"
+    camera_left_key: str = "camera_left"
+    camera_right_key: str = "camera_right"
+    # Serving prompt controls. If the client request has no instruction,
+    # default_instruction is used. prompt_template follows train-config semantics.
+    default_instruction: str | None = None
+    prompt_template: str | None = None
+    qwen25_prompt_template: str | None = None
+    prompt_priority_order: str | None = None
+
+@dataclasses.dataclass
+class Args:
+    """Arguments for the serve_wall_x script."""
+
+    # Environment mode (used for default configurations)
+    env: EnvMode = EnvMode.X2ROBOT
+
+    # Model configuration. If not provided, uses default config for the environment
+    model_config: ServerModelConfig | None = None
+
+    # Default text prompt to use if not provided in observation
+    default_prompt: str | None = None
+
+    # Port to serve the policy on
+    port: int = 43007
+
+    # Host to bind the server to
+    host: str = "0.0.0.0"
+
+    # Enable debug logging
+    debug: bool = False
+
+    # Image passing mode
+    image_passing_mode: str = "base64"  # ["numpy", "base64"]
+
+    # Model type
+    model_type: str = "wallx"  # OSS supports qwen2.5 Wall-X only
+
+    # Serialize actions via robot_preprocessor (True for robot control, False for raw output)
+    serialize_actions: bool = True
+
+    # -- Dynamic batching -----------------------------------------
+    # Set max_batch_size to enable dynamic batching. None = single mode.
+    max_batch_size: int | None = None
+    max_wait_time_ms: float = 0
+    max_queue_size: int = 100
+    timeout_ms: float = 30000
+
+    # -- Engine flags ---------------------------------------------
+    enable_experimental_engine: bool = False
+    enable_cuda_graph: bool = False
+
+
+# Default model configurations for each environment
+DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = {
+    EnvMode.X2ROBOT: ServerModelConfig(
+        checkpoint_path=None,
+        train_config_path=None,
+        robot_action_start_ratio=0.0,
+        robot_action_end_ratio=1.0,
+        robot_action_interpolate_multiplier=10,
+        robot_use_joint_angle_control=False,
+        turtle_as_desktop=False,
+        action_horizon=32,
+        action_dim=None,
+        model_device="cuda",
+        num_inference_timesteps=10,
+    ),
+    EnvMode.LIBERO: ServerModelConfig(
+        checkpoint_path=None,
+        train_config_path=None,
+        robot_type="desktop",
+        robot_action_start_ratio=0.0,
+        robot_action_end_ratio=1.0,
+        robot_action_interpolate_multiplier=1,
+        action_horizon=10,
+        action_dim=None,
+        model_device="cuda",
+        num_inference_timesteps=10,
+        cam_names=None,
+    ),
+}
+
+
+def get_model_config(args: Args) -> ServerModelConfig:
+    """Get model configuration from args or defaults."""
+    if args.model_config is not None:
+        return args.model_config
+
+    if config := DEFAULT_CONFIGS.get(args.env):
+        logger.info(f"Using default configuration for {args.env.value}")
+        return config
+
+    raise ValueError(
+        f"No default configuration for {args.env.value}. "
+        f"Please provide --model-config with model_path and action_tokenizer_path."
+    )
+
+
+def create_policy(args: Args):
+    """Create a policy from the given arguments."""
+    model_config = get_model_config(args)
+    if args.model_type != "wallx":
+        raise ValueError(
+            f"Unsupported model type: {args.model_type!r}. "
+            "The public package only supports model_type='wallx'."
+        )
+    policy = get_wallx_policy(
+        model_config, args.image_passing_mode, args.serialize_actions
+    )
+    return policy
+
+
+def main(args: Args) -> None:
+    """Main function to start the Wall-X model server."""
+    log_level = logging.DEBUG if args.debug else logging.INFO
+    logging.basicConfig(
+        level=log_level,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+
+    # Set engine environment variables
+    if args.enable_experimental_engine:
+        os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true"
+        logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true")
+
+    if args.enable_cuda_graph:
+        os.environ["ENABLE_CUDA_GRAPH"] = "true"
+        logger.info("ENABLE_CUDA_GRAPH=true")
+
+    logger.info("Starting model server")
+    logger.info(f"Model type: {args.model_type}")
+    logger.info(f"Environment: {args.env.value}")
+    logger.info(f"Port: {args.port}")
+    logger.info(f"Host: {args.host}")
+    logger.info(f"Serialize actions: {args.serialize_actions}")
+
+    # Create policy
+    try:
+        policy = create_policy(args)
+    except Exception as e:
+        logger.error(f"Failed to create policy: {e}")
+        logger.error(traceback.format_exc())
+        sys.exit(1)
+
+    # Get policy metadata
+    policy_metadata = policy.metadata
+    policy_metadata["env"] = args.env.value
+
+    # Get network info
+    hostname = socket.gethostname()
+    try:
+        local_ip = socket.gethostbyname(hostname)
+    except Exception:
+        local_ip = "unknown"
+
+    logger.info(f"Server hostname: {hostname}")
+    logger.info(f"Server IP: {local_ip}")
+    logger.info(f"Server will be available at: ws://{args.host}:{args.port}")
+    logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz")
+
+    batching_str = (
+        f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms"
+        if args.max_batch_size
+        else "disabled"
+    )
+    logger.info(f"Batching: {batching_str}")
+
+    # Create and start server
+    server = WebsocketPolicyServer(
+        policy=policy,
+        host=args.host,
+        port=args.port,
+        metadata=policy_metadata,
+        max_batch_size=args.max_batch_size,
+        max_wait_time_ms=args.max_wait_time_ms,
+        max_queue_size=args.max_queue_size,
+        timeout_ms=args.timeout_ms,
+    )
+
+    logger.info("Starting server...")
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        logger.info("Server stopped by user")
+    except Exception as e:
+        logger.error(f"Server error: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main(tyro.cli(Args))

wall_x/_vendor/harrix/serving/policy/__init__.py

@@ -0,0 +1,5 @@
+"""Wall-X policy for harrix websocket serving."""
+
+from .wall_x_policy import WallXPolicy
+
+__all__ = ["WallXPolicy"]

wall_x/_vendor/harrix/serving/policy/_smoothing.py

@@ -0,0 +1,104 @@
+"""Test-time Laplacian smoothing for serving model outputs.
+
+Mirrors the open-loop implementation in
+`wallx_bus2604/wall-x/run_scripts/infer_openloop.py`
+(_laplacian_smooth @ 203-209, smooth_action / smooth_gripper fields @ 186-187,
+application @ 458-463). Invoked by WallXPolicy before downstream 6D->euler
+conversion so that smoothed rotations propagate through serialization.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Sequence
+
+import numpy as np
+import torch
+
+from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import dof_dims
+
+
+def _laplacian_smooth(a: np.ndarray, lam: float = 1.0, iters: int = 30) -> np.ndarray:
+    """Iterative Laplacian smoothing along axis 0; endpoints pinned."""
+    a = a.copy()
+    orig = a.copy()
+    for _ in range(iters):
+        a[1:-1] = (orig[1:-1] + lam * (a[:-2] + a[2:])) / (1 + 2 * lam)
+    return a
+
+
+def _gripper_column_indices(
+    predict_action_keys: Sequence[str],
+    action_padding_dof: Optional[int] = None,
+) -> List[int]:
+    """Column indices in the flat (T, D) layout that correspond to gripper dofs."""
+    cols: List[int] = []
+    dof_start = 0
+    for key in predict_action_keys:
+        if key == "action_padding":
+            dof_start += action_padding_dof or 0
+            continue
+        short = key.replace("follow_", "").replace("master_", "")
+        width = dof_dims[short]
+        if "gripper" in short:
+            cols.extend(range(dof_start, dof_start + width))
+        dof_start += width
+    return cols
+
+
+_LAZY_ACTION_KEYS = (
+    "action_left_ee_cartesian_pos",
+    "action_right_ee_cartesian_pos",
+    "action_left_ee_rotation",
+    "action_right_ee_rotation",
+    "action_left_ee_rotation_6D",
+    "action_right_ee_rotation_6D",
+)
+
+
+def apply_smoothing(
+    model_output: dict,
+    smooth_action: bool,
+    smooth_gripper: bool,
+    predict_action_keys: Sequence[str],
+    action_padding_dof: Optional[int] = None,
+    action_dim: Optional[int] = None,
+) -> None:
+    """Smooth `model_output['predict_action']` in place and refresh per-arm keys."""
+    if not smooth_action:
+        return
+    pa = model_output.get("predict_action")
+    if pa is None:
+        return
+
+    was_tensor = isinstance(pa, torch.Tensor)
+    arr = pa.detach().cpu().numpy() if was_tensor else np.asarray(pa)
+
+    orig_ndim = arr.ndim
+    if orig_ndim == 3:
+        if arr.shape[0] != 1:
+            return
+        arr = arr[0]
+    if arr.ndim != 2 or arr.shape[0] < 3:
+        return
+    if action_dim is not None and arr.shape[-1] != action_dim:
+        return
+
+    orig = arr.copy()
+    smoothed = _laplacian_smooth(arr)
+
+    if not smooth_gripper:
+        for c in _gripper_column_indices(predict_action_keys, action_padding_dof):
+            if c < smoothed.shape[-1]:
+                smoothed[:, c] = orig[:, c]
+
+    out = smoothed[None] if orig_ndim == 3 else smoothed
+    if was_tensor:
+        out = torch.from_numpy(out).to(device=pa.device, dtype=pa.dtype)
+    model_output["predict_action"] = out
+
+    rsd = model_output.get("robot_state_action_data")
+    if rsd is not None:
+        for k in _LAZY_ACTION_KEYS:
+            if k in rsd.data:
+                rsd.data[k] = None
+        rsd.save_action_data(out)

wall_x/_vendor/harrix/serving/policy/wall_x_policy.py

@@ -0,0 +1,328 @@
+import base64
+import logging
+from typing import Dict, Any, List
+
+import cv2
+import numpy as np
+import torch
+from wall_x._vendor.harrix.serving.websocket_policy_server import BasePolicy
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper import WallxModelWrapper
+from wall_x._vendor.harrix.serving._wallx_infer.robot import (
+    DesktopRobotPreprocessor,
+    TurtleRobotPreprocessor,
+    EX001RobotPreprocessor,
+)
+from wall_x._vendor.harrix.serving.policy._smoothing import apply_smoothing
+from wall_x.utils.timers import ScopeTimer
+
+logger = logging.getLogger(__name__)
+
+# Inference modes aligned with model_wrapper for vqa / batch extensions
+INFER_MODE_FLOW = "flow"
+INFER_MODE_AR = "ar"
+INFER_MODE_FLOW_WITH_SUBTASK = "flow_with_subtask"
+INFER_MODE_DLLM_FLOW = "dllm"
+INFER_MODE_DLLM_DD = "discrete_diffusion"
+ACTION_INFER_MODES = (
+    INFER_MODE_FLOW,
+    INFER_MODE_AR,
+    INFER_MODE_FLOW_WITH_SUBTASK,
+    INFER_MODE_DLLM_FLOW,
+    INFER_MODE_DLLM_DD,
+)
+
+
+class WallXPolicy(BasePolicy):
+    """Policy wrapper for Wall-X model that implements the BasePolicy interface."""
+
+    def __init__(
+        self,
+        config: InferConfig,
+        image_passing_mode: str = "base64",
+        default_infer_mode: str = INFER_MODE_FLOW,
+        serialize_actions: bool = True,
+    ):
+        """Initialize the Wall-X policy.
+
+        Args:
+            config: Inference configuration dataclass.
+            image_passing_mode: How images are passed from client ('base64' or 'numpy').
+            default_infer_mode: Default inference mode ('flow', 'ar', 'flow_with_subtask', etc.).
+            serialize_actions: If True, actions are serialized via robot_preprocessor;
+                if False, raw model output is returned directly.
+        """
+        self.config = config
+        self.model_wrapper = WallxModelWrapper(config)
+        self.robot_preprocessor = self._register_robot_preprocessor()
+        self.image_passing_mode = image_passing_mode
+        self.default_infer_mode = default_infer_mode
+        self.serialize_actions = serialize_actions
+        logger.info(
+            "Image passing mode: %s, robot_type: %s, default_infer_mode: %s, "
+            "serialize_actions: %s, smooth_action: %s, smooth_gripper: %s",
+            self.image_passing_mode,
+            config.robot_type,
+            self.default_infer_mode,
+            self.serialize_actions,
+            config.smooth_action,
+            config.smooth_gripper,
+        )
+
+    def _register_robot_preprocessor(self):
+        """Select preprocessor by config.robot_type; mirrors env._register_robot."""
+        if self.config.robot_type == "desktop":
+            return DesktopRobotPreprocessor(self.config)
+        if self.config.robot_type == "turtle":
+            return TurtleRobotPreprocessor(self.config)
+        if self.config.robot_type == "ex001":
+            return EX001RobotPreprocessor(self.config)
+        raise ValueError(f"Invalid robot_type: {self.config.robot_type!r}")
+
+    @staticmethod
+    def _decode_base64_image(image_b64: str) -> np.ndarray:
+        """Decode client base64 JPEG/PNG payloads to RGB images."""
+        img_bytes = base64.b64decode(image_b64)
+        img_array = np.frombuffer(img_bytes, dtype=np.uint8)
+        decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+        if decoded_img is None:
+            raise ValueError("Failed to decode base64 image payload")
+        return cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB)
+
+    def reset(self) -> None:
+        """Reset the policy state."""
+        self.action_buffer = []
+        self.buffer_index = 0
+        logger.debug("Policy reset")
+
+    def _get_dof_config(self) -> dict:
+        train_config = self.config.train_config or {}
+        return (
+            train_config.get("dof_config")
+            or train_config.get("task", {}).get("dof_config", {})
+            or {}
+        )
+
+    def _is_single_arm_right_only(self) -> bool:
+        """True for single-arm LIBERO-style configs (right arm only, no left arm)."""
+        train_config = self.config.train_config or {}
+        agent_pos = (
+            train_config.get("agent_pos_config")
+            or train_config.get("task", {}).get("agent_pos_config")
+            or {}
+        )
+        skip = {"action_padding"}
+        has_left = any(k.startswith("follow_left_") for k in agent_pos if k not in skip)
+        has_right = any(
+            k.startswith("follow_right_") for k in agent_pos if k not in skip
+        )
+        return has_right and not has_left
+
+    def _pack_action_chunk_response(
+        self, model_output: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Return a msgpack-safe action chunk for websocket clients."""
+        predict_action = model_output["predict_action"]
+        if isinstance(predict_action, torch.Tensor):
+            predict_action = predict_action.detach().cpu().numpy()
+        response: Dict[str, Any] = {
+            "predict_action": np.asarray(predict_action, dtype=np.float32),
+        }
+        if "subtask" in model_output:
+            response["subtask"] = model_output["subtask"]
+        return response
+
+    def _get_predict_action_keys(self) -> List[str]:
+        """Resolve flat action key order for smoothing / serialization."""
+        data_cfg = self.config.data_config
+        keys = None
+        if isinstance(data_cfg, dict):
+            keys = data_cfg.get("predict_action_keys")
+        else:
+            keys = getattr(data_cfg, "predict_action_keys", None)
+        if keys:
+            return list(keys)
+        return list(self._get_dof_config().keys())
+
+    def _apply_smoothing(self, model_output: Dict[str, Any]) -> None:
+        """Apply Laplacian smoothing on predict_action before 6D->euler conversion."""
+        cfg = self.config
+        if not getattr(cfg, "smooth_action", False):
+            return
+        dof_config = self._get_dof_config()
+        apply_smoothing(
+            model_output,
+            smooth_action=True,
+            smooth_gripper=cfg.smooth_gripper,
+            predict_action_keys=self._get_predict_action_keys(),
+            action_padding_dof=dof_config.get("action_padding"),
+            action_dim=cfg.action_dim,
+        )
+
+    def _run_action_infer(
+        self, observation: Dict, instruction: str, mode: str
+    ) -> Dict[str, Any]:
+        """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data.
+
+        Supports flow | ar | flow_with_subtask. VQA can be added here later (different return format).
+        """
+        if mode == INFER_MODE_FLOW:
+            return self.model_wrapper.infer_flow_action(observation, instruction)
+        if mode == INFER_MODE_AR:
+            return self.model_wrapper.infer_ar_action(observation, instruction)
+        if mode == INFER_MODE_FLOW_WITH_SUBTASK:
+            with ScopeTimer("infer_subtask"):
+                subtask = self.model_wrapper.infer_subtask(observation, instruction)
+            with ScopeTimer("infer_flow_action"):
+                model_output = self.model_wrapper.infer_flow_action(
+                    observation, subtask
+                )
+            model_output["subtask"] = subtask
+            return model_output
+        if mode == INFER_MODE_DLLM_FLOW:
+            return self.model_wrapper.infer_dllm_action(
+                observation, instruction, use_ar_action=False
+            )
+        if mode == INFER_MODE_DLLM_DD:
+            return self.model_wrapper.infer_dllm_action(
+                observation, instruction, use_ar_action=True
+            )
+        raise ValueError(
+            f"Unsupported infer_mode={mode!r}, expected one of {ACTION_INFER_MODES}"
+        )
+
+    def infer(self, obs: Dict) -> Dict:
+        """Infer action from observation.
+
+        Args:
+            obs: Dictionary containing:
+                - 'state': Robot state
+                - 'views': Camera views (keyed by camera name)
+                - 'instruction': Task instruction
+                - Optional: 'infer_mode' - one of 'flow' | 'ar' | 'flow_with_subtask'
+                - Optional: 'robot_action_start_ratio' / 'robot_action_end_ratio' /
+                  'robot_action_interpolate_multiplier' - override config for action trim/interpolate
+
+        Returns:
+            When serialize_actions=True (default): Serialized action dict
+                (e.g. follow1_pos/follow2_pos or follow1_joints/follow2_joints).
+            When serialize_actions=False: Raw model_output dict.
+            If infer_mode is flow_with_subtask, also includes 'subtask'.
+        """
+        state = obs["state"]
+        views = obs["views"]
+        instruction = obs.get("instruction") or self.config.default_instruction or ""
+        if self.image_passing_mode == "base64":
+            for k, v in views.items():
+                views[k] = np.expand_dims(self._decode_base64_image(v), axis=0)
+
+        with ScopeTimer("get_observation"):
+            observation = self.robot_preprocessor.get_observation(state, views)
+
+        mode = obs.get("infer_mode", self.default_infer_mode)
+        with ScopeTimer(f"infer_{mode}"):
+            model_output = self._run_action_infer(observation, instruction, mode)
+
+        self._apply_smoothing(model_output)
+        if (
+            "robot_state_action_data" in model_output
+            and "predict_action" in model_output
+        ):
+            model_output["robot_state_action_data"].save_action_data(
+                model_output["predict_action"]
+            )
+
+        # Single-arm LIBERO has no left-arm EE to serialize; dual-arm follow1/follow2
+        # layout would fail in get_serialized_actions. Return raw action chunks instead.
+        if not self.serialize_actions or self._is_single_arm_right_only():
+            return self._pack_action_chunk_response(model_output)
+
+        return self.robot_preprocessor.get_serialized_actions(
+            model_output, robot_action_interpolate_multiplier=1
+        )  # interpolation on robot websocket client
+
+    # -- Batch inference ----------------------------------------------
+
+    def _preprocess_obs(self, obs: Dict):
+        """Preprocess a single observation dict into (observation, instruction).
+
+        Handles both base64 and raw image modes.
+        """
+        state = obs["state"]
+        views = obs["views"]
+        instruction = obs.get("instruction") or self.config.default_instruction or ""
+
+        if self.image_passing_mode == "base64":
+            for k, v in views.items():
+                views[k] = np.expand_dims(self._decode_base64_image(v), axis=0)
+        else:
+            # raw numpy: ensure (1, H, W, C)
+            for k, img in views.items():
+                if isinstance(img, np.ndarray) and img.ndim == 3:
+                    views[k] = np.expand_dims(img, axis=0)
+
+        observation = self.robot_preprocessor.get_observation(state, views)
+        return observation, instruction
+
+    def infer_batch(
+        self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False
+    ) -> List[Dict[str, Any]]:
+        """Perform batch inference.
+
+        Args:
+            obs_list: List of observations, each containing:
+                - "views": dict of camera images
+                - "state": robot state dict or array
+                - "instruction": text instruction
+
+        Returns:
+            List of action dicts.
+        """
+        batch_size = len(obs_list)
+        logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations")
+
+        try:
+            observations = []
+            instructions = []
+            for obs in obs_list:
+                observation, instruction = self._preprocess_obs(obs)
+                observations.append(observation)
+                instructions.append(instruction)
+
+            mode = obs_list[0].get("infer_mode", self.default_infer_mode)
+
+            with torch.no_grad():
+                if mode == INFER_MODE_FLOW:
+                    model_outputs = self.model_wrapper.infer_flow_action_batch(
+                        observations, instructions
+                    )
+                else:
+                    # AR / flow_with_subtask: fall back to per-sample inference
+                    model_outputs = []
+                    for obs_dict, instruction in zip(observations, instructions):
+                        output = self._run_action_infer(obs_dict, instruction, mode)
+                        model_outputs.append(output)
+
+            if skip_serialize:
+                return [{}] * len(model_outputs)
+
+            results = []
+            for model_output in model_outputs:
+                self._apply_smoothing(model_output)
+                if self.serialize_actions and not self._is_single_arm_right_only():
+                    action = self.robot_preprocessor.get_serialized_actions(
+                        model_output, robot_action_interpolate_multiplier=1
+                    )
+                    results.append(action)
+                else:
+                    results.append(self._pack_action_chunk_response(model_output))
+
+            return results
+
+        except Exception as e:
+            logger.error(f"Batch inference failed: {e}", exc_info=True)
+            return [{"action": {}, "error": str(e)} for _ in obs_list]
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        return {"batch_enabled": True, "model": "wall-x"}

wall_x/_vendor/harrix/serving/scheduler.py

@@ -0,0 +1,160 @@
+import asyncio
+import logging
+import time
+import uuid
+from collections import deque
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Request:
+    request_id: str
+    obs: Dict[str, Any]
+    future: asyncio.Future
+    timestamp: float
+
+
+class RequestScheduler:
+    def __init__(
+        self,
+        policy,
+        max_batch_size: int = 8,
+        max_wait_time_ms: float = 100,
+        max_queue_size: int = 128,
+        timeout_ms: float = 5000,
+    ):
+        self.policy = policy
+        self.max_batch_size = max_batch_size
+        self.max_wait_time = max_wait_time_ms / 1000.0
+        self.max_queue_size = max_queue_size
+        self.timeout = timeout_ms / 1000.0
+
+        self.queue = deque()
+        self.queue_lock = asyncio.Lock()
+        self.queue_not_empty = asyncio.Condition(self.queue_lock)
+
+        self.running = False
+        self.batch_task = None
+
+    async def start(self):
+        self.running = True
+        self.batch_task = asyncio.create_task(self._batch_loop())
+        logger.info(
+            f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms"
+        )
+
+    async def stop(self):
+        self.running = False
+        # Process all remaining requests before shutting down
+        async with self.queue_lock:
+            remaining = len(self.queue)
+        if remaining:
+            logger.info(
+                f"Graceful shutdown: processing {remaining} remaining request(s)"
+            )
+        while True:
+            async with self.queue_lock:
+                if len(self.queue) == 0:
+                    break
+                batch = []
+                while self.queue and len(batch) < self.max_batch_size:
+                    batch.append(self.queue.popleft())
+            if batch:
+                await self._process_batch(batch)
+        # Now stop the batch loop
+        async with self.queue_lock:
+            self.queue_not_empty.notify_all()
+        if self.batch_task:
+            await self.batch_task
+        logger.info("RequestScheduler stopped")
+
+    async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]:
+        request_id = str(uuid.uuid4())
+        future = asyncio.Future()
+        request = Request(
+            request_id=request_id, obs=obs, future=future, timestamp=time.monotonic()
+        )
+
+        async with self.queue_lock:
+            if len(self.queue) >= self.max_queue_size:
+                raise RuntimeError(
+                    f"Queue full: {len(self.queue)}/{self.max_queue_size}"
+                )
+            self.queue.append(request)
+            self.queue_not_empty.notify()
+
+        try:
+            result = await asyncio.wait_for(future, timeout=self.timeout)
+            return result
+        except asyncio.TimeoutError:
+            logger.error(f"Request {request_id} timeout after {self.timeout}s")
+            raise
+
+    async def _batch_loop(self):
+        while self.running:
+            try:
+                batch = await self._collect_batch()
+                if batch:
+                    await self._process_batch(batch)
+            except Exception as e:
+                logger.error(f"Unexpected error in batch loop: {e}", exc_info=True)
+
+    async def _collect_batch(self) -> List[Request]:
+        async with self.queue_not_empty:
+            while self.running and len(self.queue) == 0:
+                await self.queue_not_empty.wait()
+
+            if not self.running:
+                return []
+
+            batch = []
+            deadline = time.monotonic() + self.max_wait_time
+
+            while len(batch) < self.max_batch_size:
+                if len(self.queue) > 0:
+                    batch.append(self.queue.popleft())
+
+                if len(batch) >= self.max_batch_size:
+                    break
+
+                if len(self.queue) == 0:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        break
+                    try:
+                        await asyncio.wait_for(
+                            self.queue_not_empty.wait(), timeout=remaining
+                        )
+                    except asyncio.TimeoutError:
+                        break
+
+            return batch
+
+    async def _process_batch(self, batch: List[Request]):
+        if not batch:
+            return
+
+        start_time = time.monotonic()
+        logger.info(f"Processing batch of {len(batch)} requests")
+
+        try:
+            obs_list = [req.obs for req in batch]
+            results = await asyncio.to_thread(self.policy.infer_batch, obs_list)
+
+            for req, result in zip(batch, results):
+                if not req.future.done():
+                    req.future.set_result(result)
+
+            infer_time = time.monotonic() - start_time
+            logger.info(
+                f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s"
+            )
+
+        except Exception as e:
+            logger.error(f"Batch processing failed: {e}", exc_info=True)
+            for req in batch:
+                if not req.future.done():
+                    req.future.set_exception(e)

wall_x/_vendor/harrix/serving/websocket_policy_server.py

@@ -0,0 +1,184 @@
+import asyncio
+import http
+import logging
+import time
+import traceback
+from typing import Dict, Any, List
+
+try:
+    import msgpack
+    import msgpack_numpy as m
+
+    m.patch()
+except ImportError:
+    logging.warning(
+        "msgpack-numpy not installed. Install with: pip install msgpack-numpy"
+    )
+    msgpack = None
+
+import websockets.asyncio.server as _server
+import websockets.frames
+
+logger = logging.getLogger(__name__)
+
+
+class BasePolicy:
+    """Base class for policies that can be served."""
+
+    def infer(self, obs: Dict) -> Dict:
+        """Infer actions from observations."""
+        raise NotImplementedError
+
+    def infer_batch(self, obs_list: List[Dict]) -> List[Dict]:
+        """Batch inference. Default implementation calls infer() per sample."""
+        return [self.infer(obs) for obs in obs_list]
+
+    def reset(self) -> None:
+        """Reset the policy to its initial state."""
+        pass
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """Return metadata about the policy."""
+        return {}
+
+
+class WebsocketPolicyServer:
+    """Serves a policy using the websocket protocol.
+
+    Implements a websocket server that:
+    1. Sends policy metadata on connection
+    2. Receives observations
+    3. Returns predicted actions (single or batched)
+    4. Tracks timing information
+
+    When batching parameters are provided, dynamic batching is enabled:
+    requests from multiple clients are queued and processed in batches.
+    """
+
+    def __init__(
+        self,
+        policy: BasePolicy,
+        host: str = "0.0.0.0",
+        port: int = 8000,
+        metadata: Dict | None = None,
+        # Dynamic batching parameters (None = disabled)
+        max_batch_size: int | None = None,
+        max_wait_time_ms: float | None = None,
+        max_queue_size: int = 100,
+        timeout_ms: float = 30000,
+    ) -> None:
+        self._policy = policy
+        self._host = host
+        self._port = port
+        self._metadata = metadata or {}
+
+        # Dynamic batching
+        self._scheduler = None
+        if max_batch_size is not None:
+            from .scheduler import RequestScheduler
+
+            self._scheduler = RequestScheduler(
+                policy=policy,
+                max_batch_size=max_batch_size,
+                max_wait_time_ms=(
+                    max_wait_time_ms if max_wait_time_ms is not None else 0
+                ),
+                max_queue_size=max_queue_size,
+                timeout_ms=timeout_ms,
+            )
+
+        logging.getLogger("websockets.server").setLevel(logging.INFO)
+
+    @property
+    def batching_enabled(self) -> bool:
+        return self._scheduler is not None
+
+    def serve_forever(self) -> None:
+        asyncio.run(self.run())
+
+    async def run(self):
+        # Start the scheduler if batching is enabled
+        if self._scheduler is not None:
+            await self._scheduler.start()
+
+        try:
+            async with _server.serve(
+                self._handler,
+                self._host,
+                self._port,
+                compression=None,
+                max_size=None,
+                ping_interval=None,
+                ping_timeout=None,
+                process_request=_health_check,
+            ) as server:
+                mode_str = "batched" if self.batching_enabled else "single"
+                logger.info(
+                    f"Server started on {self._host}:{self._port} (mode={mode_str})"
+                )
+                await server.serve_forever()
+        finally:
+            if self._scheduler is not None:
+                await self._scheduler.stop()
+
+    async def _handler(self, websocket: _server.ServerConnection):
+        logger.info(f"Connection from {websocket.remote_address} opened")
+
+        if msgpack is None:
+            await websocket.close(
+                code=websockets.frames.CloseCode.INTERNAL_ERROR,
+                reason="msgpack-numpy not installed on server",
+            )
+            return
+
+        # Send metadata to client
+        await websocket.send(msgpack.packb(self._metadata))
+
+        prev_total_time = None
+        while True:
+            try:
+                start_time = time.monotonic()
+                obs = msgpack.unpackb(await websocket.recv())
+
+                infer_time = time.monotonic()
+
+                if self._scheduler is not None:
+                    # Dynamic batching path
+                    action = await self._scheduler.add_request(obs)
+                else:
+                    # Single inference path
+                    action = await asyncio.to_thread(self._policy.infer, obs)
+
+                infer_time = time.monotonic() - infer_time
+
+                action["server_timing"] = {
+                    "infer_ms": infer_time * 1000,
+                }
+                if prev_total_time is not None:
+                    action["server_timing"]["prev_total_ms"] = prev_total_time * 1000
+
+                await websocket.send(msgpack.packb(action))
+                prev_total_time = time.monotonic() - start_time
+
+            except websockets.ConnectionClosed:
+                logger.info(f"Connection from {websocket.remote_address} closed")
+                break
+            except Exception as e:
+                logger.error(f"Error handling request: {e}")
+                await websocket.send(traceback.format_exc())
+                await websocket.close(
+                    code=websockets.frames.CloseCode.INTERNAL_ERROR,
+                    reason="Internal server error. Traceback included in previous frame.",
+                )
+                raise
+
+
+def _health_check(
+    connection: _server.ServerConnection, request: _server.Request
+) -> _server.Response | None:
+    if request.path == "/healthz":
+        return connection.respond(http.HTTPStatus.OK, "OK\n")
+    if request.path == "/v2/health/ready":
+        return connection.respond(http.HTTPStatus.OK, "OK\n")
+    return None