From 88e6e9bee6aa0dc079556dd8d7e6f1dcbe517fef Mon Sep 17 00:00:00 2001 From: Saberlve Date: Wed, 19 Aug 2026 09:11:32 +0000 Subject: [PATCH] Add starVLA remote policy eval for xArm7 - uf_starvla_eval.py: eval loop that queries a starVLA WebSocket policy server (e.g. pi05) instead of a local LeRobot policy; reuses robot connect/reset, keyboard control, precise_sleep pacing and the optional ActionSafetyGuard from uf_lerobot_eval.py - utils/starvla_ws_client: vendored WebsocketClientPolicy + msgpack_numpy from starVLA deployment/model_server/tools - config/eval/xarm7_starvla_eval_config.yaml: single-camera eval config - pyproject.toml: uf-starvla-eval entry point + websockets/msgpack deps --- config/eval/xarm7_starvla_eval_config.yaml | 57 +++++ pyproject.toml | 3 + .../scripts/uf_starvla_eval.py | 207 ++++++++++++++++++ .../utils/starvla_ws_client/__init__.py | 10 + .../utils/starvla_ws_client/msgpack_numpy.py | 59 +++++ .../websocket_policy_client.py | 196 +++++++++++++++++ 6 files changed, 532 insertions(+) create mode 100644 config/eval/xarm7_starvla_eval_config.yaml create mode 100644 src/lerobot_robot_ufactory/scripts/uf_starvla_eval.py create mode 100644 src/lerobot_robot_ufactory/utils/starvla_ws_client/__init__.py create mode 100644 src/lerobot_robot_ufactory/utils/starvla_ws_client/msgpack_numpy.py create mode 100644 src/lerobot_robot_ufactory/utils/starvla_ws_client/websocket_policy_client.py diff --git a/config/eval/xarm7_starvla_eval_config.yaml b/config/eval/xarm7_starvla_eval_config.yaml new file mode 100644 index 0000000..5e0a2e6 --- /dev/null +++ b/config/eval/xarm7_starvla_eval_config.yaml @@ -0,0 +1,57 @@ +# Real-robot eval config for uf-starvla-eval (external starVLA policy server). +# Robot section is derived from config/gello/xarm7_gello_record_config.yaml +# (teleop / dataset / web_preview sections removed). +robot: + type: uf::robot + id: "uf_robot" + robot_dof: 7 + control_space: "joint" + # TODO: confirm this is your xArm controller IP. + robot_ip: "192.168.1.245" + # 2: xArm Gripper G2 (0-84 mm opening range). + gripper_type: 2 + enable_logs: false + # Up to 60 Hz goal updates; unchanged targets are filtered below. + gripper_command_interval_s: 0.0166667 + # xArm Gripper G2 speed, 15-225 mm/s; 100 is the SDK default. + gripper_speed: 100 + # xArm Gripper G2 gripping force, 1-100 percent; 50 is the SDK default. + gripper_force: 50 + # Use the high-frequency servo interface for lower-latency action streaming. + joint_command_mode: 1 + max_joint_velocity: 120 + # TCP z floor in the xArm base coordinate system (mm). + min_tcp_z_mm: -2.0 + # CPU-local FK/Jacobian projection keeps ServoJ free of synchronous SDK queries. + tcp_z_guard_backend: "local_projection" + tcp_z_soft_margin_mm: 0.5 + local_kinematics_max_error_mm: 2.0 + controller_safety_boundary: true + # Append gripper initialization/read/write failures here. + gripper_error_log_path: "logs/xarm7_gripper_errors.log" + cameras: + # Single camera view; must match `camera_key` below and the training setup. + # TODO: fill in the serial number of YOUR RealSense camera. + camera: + type: intelrealsense + serial_number_or_name: "242622070583" + width: 640 + height: 480 + fps: 30 + +# starVLA policy server address (server binds 0.0.0.0; set the server IP here +# if the server runs on a different machine). +server_host: "127.0.0.1" +server_port: 10093 + +# Control frequency for streaming actions to the robot. +fps: 30 +# Execute the first N steps of each predicted action chunk (T=50), then re-infer. +# N=1 means fully closed-loop (re-infer every step). +steps_per_inference: 25 +single_task: "Pick up the black bottle and place it on the blue bag" +n_episodes: 50 +# Key of the camera in the robot observation dict (camera name above). +camera_key: "camera" +# Enable the action safety guard (thresholds in ActionSafetyConfig). +enable_safety: false diff --git a/pyproject.toml b/pyproject.toml index c2cac1e..834eaca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,8 @@ dependencies = [ "lerobot[intelrealsense]==0.4.3", "xarm-python-sdk", "opencv-python", + "websockets", + "msgpack", ] [project.scripts] @@ -37,6 +39,7 @@ uf-robot-teleop = "lerobot_robot_ufactory.scripts.uf_robot_teleop:main" uf-lerobot-record = "lerobot_robot_ufactory.scripts.uf_lerobot_record:main" record = "lerobot_robot_ufactory.scripts.uf_lerobot_record:main" uf-lerobot-eval = "lerobot_robot_ufactory.scripts.uf_lerobot_eval:main" +uf-starvla-eval = "lerobot_robot_ufactory.scripts.uf_starvla_eval:main" uf-lerobot-replay = "lerobot_robot_ufactory.scripts.uf_lerobot_replay:main" replay = "lerobot_robot_ufactory.scripts.uf_lerobot_replay:main" uf-vive-calibrate = "lerobot_robot_ufactory.scripts.vive_calibrate:main" diff --git a/src/lerobot_robot_ufactory/scripts/uf_starvla_eval.py b/src/lerobot_robot_ufactory/scripts/uf_starvla_eval.py new file mode 100644 index 0000000..2204731 --- /dev/null +++ b/src/lerobot_robot_ufactory/scripts/uf_starvla_eval.py @@ -0,0 +1,207 @@ +"""Real-robot evaluation against an external starVLA policy server. + +Usage: + # 1. Start the starVLA policy server (in the starVLA repo / environment): + python deployment/model_server/server_policy.py \ + --ckpt_path --port 10093 --use_bf16 + + # 2. Run this eval script on the robot machine: + uf-starvla-eval --config_path config/eval/xarm7_starvla_eval_config.yaml + +The policy runs on the server; this script only streams observations +(joint state + one camera image + task text) over WebSocket and executes the +returned action chunk on the xArm. + +Keyboard controls (same as uf_lerobot_eval): + Right/Left arrow : reset current episode (robot returns to initial pose) + Esc : exit eval loop and disconnect +""" + +import logging +import time +from dataclasses import asdict, dataclass +from pprint import pformat + +import numpy as np + +import lerobot_robot_ufactory # patch: registers uf:: robot/camera types +# Register camera config subclasses ("opencv", "intelrealsense") so draccus +# can decode the robot.cameras section; uf_lerobot_eval gets these transitively +# via lerobot.scripts.lerobot_record, which this script does not import. +from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.configs import parser +from lerobot.robots import ( # noqa: F401 + Robot, + RobotConfig, + make_robot_from_config, +) +from lerobot.utils.control_utils import is_headless +from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.robot_utils import precise_sleep +from lerobot.utils.utils import init_logging + +from lerobot_robot_ufactory.utils.action_safety import ActionSafetyConfig, ActionSafetyGuard +from lerobot_robot_ufactory.utils.starvla_ws_client import WebsocketClientPolicy +from lerobot_robot_ufactory.utils.utils import init_keyboard_listener + + +@dataclass +class StarVLAEvalConfig: + robot: RobotConfig + # starVLA policy server address (server binds 0.0.0.0; use the server IP here) + server_host: str = "127.0.0.1" + server_port: int = 10093 + # Control frequency for streaming actions to the robot. + fps: int = 30 + # Execute the first N steps of each predicted action chunk, then re-infer. + # N=1 means fully closed-loop (re-infer every step). Server chunks are T=50. + steps_per_inference: int = 25 + single_task: str = "Pick up the black bottle and place it on the blue bag" + n_episodes: int = 50 + # Key of the camera in the robot observation dict (camera name in robot config). + camera_key: str = "camera" + # Enable the action safety guard (thresholds in ActionSafetyConfig). + enable_safety: bool = False + + +def _build_state(obs: dict) -> np.ndarray: + """8-dim proprio state: 7 joint positions (rad) + gripper (0=open, 1=close).""" + return np.array([obs[f"J{i}.pos"] for i in range(1, 8)] + [obs["gripper.pos"]], dtype=np.float32) + + +def _build_action_dict(action: np.ndarray) -> dict: + """Map one (8,) action row to the robot action dict. + + Actions are already denormalized by the server: absolute joint positions + (rad) + gripper in [0, 1] (0=open, 1=close), matching robot conventions. + """ + action_dict = {f"J{i + 1}.pos": float(action[i]) for i in range(7)} + action_dict["gripper.pos"] = float(action[7]) + return action_dict + + +def eval_loop(cfg: StarVLAEvalConfig, safety_guard: ActionSafetyGuard | None = None): + init_logging() + logging.info(pformat(asdict(cfg))) + if safety_guard is not None: + safety_guard.log_config() + + robot = make_robot_from_config(cfg.robot) + robot.connect() + + client = WebsocketClientPolicy(cfg.server_host, cfg.server_port) + # Echoed so the operator can manually verify action_chunk_size, image + # size/count etc. against the training setup before running episodes. + logging.info(f"starVLA server metadata: {pformat(client.get_server_metadata())}") + + events = {"reset": False, "exit": False} + listener = None + + if not is_headless(): + from pynput import keyboard + + def on_press(key): + try: + if key == keyboard.Key.right: + print("Right arrow key pressed. Resetting...") + events["reset"] = True + elif key == keyboard.Key.left: + print("Left arrow key pressed. Resetting....") + events["reset"] = True + elif key == keyboard.Key.esc: + print("Escape key pressed. Stopping ...") + events["exit"] = True + except Exception as e: + print(f"Error handling key press: {e}") + + listener, events = init_keyboard_listener(events=events, on_press=on_press) + + sleep_time_s = 1 / cfg.fps + + print("\n********** starVLA Policy Eval Episode Loop Start **********") + + try: + episode = 0 + while episode < cfg.n_episodes and not events["exit"]: + print(f"\n********** Episode {episode + 1}/{cfg.n_episodes} **********") + reset = getattr(robot, "reset_to_initial", None) + if reset is None: + reset = robot.configure + reset() + events["reset"] = False + + while True: + if events["reset"] or events["exit"]: + events["reset"] = False + print("\n********** starVLA Policy Eval Episode (Reset) **********") + break + + # Get robot observation + obs = robot.get_observation() + state = _build_state(obs) + image = obs[cfg.camera_key] # uint8 HWC RGB + + # NOTE: inference is blocking (one flow-matching pass can take + # several hundred ms) and no actions are sent while waiting. + # This is acceptable here: xArm ServoJ holds the last commanded + # position, so the arm simply pauses between action chunks. + resp = client.predict_action( + {"examples": [{"image": [image], "lang": cfg.single_task, "state": state}]} + ) + actions = np.asarray(resp["data"]["actions"][0]) # (T, 8), denormalized + + # Execute the first N steps of the chunk, then re-infer. + for action in actions[: cfg.steps_per_inference]: + if events["reset"] or events["exit"]: + break + + start_loop_t = time.perf_counter() + action_dict = _build_action_dict(action) + + # Safety check: joint-space pose limits do not apply (the + # guard only checks TCP-pose actions and gripper NaN/Inf), + # but it still catches non-finite gripper commands. + if safety_guard is not None: + violation = safety_guard.check(action_dict, {}, [""]) + if violation is not None: + logging.error(f"*** SAFETY HALT *** {violation}") + print( + f"\n*** SAFETY HALT *** {violation}\n" + "Action was NOT sent. Press right arrow (->) to reset and resume, ESC to exit." + ) + events["reset"] = True + break + + robot.send_action(action_dict) + + dt_s = time.perf_counter() - start_loop_t + precise_sleep(sleep_time_s - dt_s) + + episode += 1 + + finally: + print("\n********** starVLA Policy Eval Loop Exit **********") + client.close() + if robot.is_connected: + robot.disconnect() + if not is_headless() and listener is not None: + listener.stop() + + +@parser.wrap() +def get_cfg(cfg: StarVLAEvalConfig) -> StarVLAEvalConfig: + return cfg + + +def main(): + register_third_party_plugins() + cfg = get_cfg() + # Action safety guard: tune thresholds in ActionSafetyConfig directly. + # Any violation triggers an e-stop; press right arrow to resume. + safety_guard = ActionSafetyGuard(ActionSafetyConfig(enabled=cfg.enable_safety)) + eval_loop(cfg, safety_guard) + + +if __name__ == "__main__": + main() diff --git a/src/lerobot_robot_ufactory/utils/starvla_ws_client/__init__.py b/src/lerobot_robot_ufactory/utils/starvla_ws_client/__init__.py new file mode 100644 index 0000000..8b79c53 --- /dev/null +++ b/src/lerobot_robot_ufactory/utils/starvla_ws_client/__init__.py @@ -0,0 +1,10 @@ +"""Vendored starVLA WebSocket policy client. + +Copied from the starVLA repository (deployment/model_server/tools/) so the +real-robot eval script can talk to a starVLA policy server without adding a +dependency on the starVLA package itself. +""" + +from .websocket_policy_client import WebsocketClientPolicy + +__all__ = ["WebsocketClientPolicy"] diff --git a/src/lerobot_robot_ufactory/utils/starvla_ws_client/msgpack_numpy.py b/src/lerobot_robot_ufactory/utils/starvla_ws_client/msgpack_numpy.py new file mode 100644 index 0000000..97782b5 --- /dev/null +++ b/src/lerobot_robot_ufactory/utils/starvla_ws_client/msgpack_numpy.py @@ -0,0 +1,59 @@ +# Vendored from starVLA: deployment/model_server/tools/msgpack_numpy.py (2026-08-19). +# Source repo: https://github.com/starVLA/starVLA +"""Adds NumPy array support to msgpack. + +msgpack is good for (de)serializing data over a network for multiple reasons: +- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution) +- msgpack is widely used and has good cross-language support +- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed + languages like Python and JavaScript +- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster + than pickle for serializing large arrays using the below strategy + +The code below is adapted from https://github.com/lebedov/msgpack-numpy. The reason not to use that library directly is +that it falls back to pickle for object arrays. +""" + +import functools + +import msgpack +import numpy as np + + +def pack_array(obj): + if (isinstance(obj, (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {obj.dtype}") + + if isinstance(obj, np.ndarray): + return { + b"__ndarray__": True, + b"data": obj.tobytes(), + b"dtype": obj.dtype.str, + b"shape": obj.shape, + } + + if isinstance(obj, np.generic): + return { + b"__npgeneric__": True, + b"data": obj.item(), + b"dtype": obj.dtype.str, + } + + return obj + + +def unpack_array(obj): + if b"__ndarray__" in obj: + return np.ndarray(buffer=obj[b"data"], dtype=np.dtype(obj[b"dtype"]), shape=obj[b"shape"]) + + if b"__npgeneric__" in obj: + return np.dtype(obj[b"dtype"]).type(obj[b"data"]) + + return obj + + +Packer = functools.partial(msgpack.Packer, default=pack_array) +packb = functools.partial(msgpack.packb, default=pack_array) + +Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array) +unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array) diff --git a/src/lerobot_robot_ufactory/utils/starvla_ws_client/websocket_policy_client.py b/src/lerobot_robot_ufactory/utils/starvla_ws_client/websocket_policy_client.py new file mode 100644 index 0000000..53d87ca --- /dev/null +++ b/src/lerobot_robot_ufactory/utils/starvla_ws_client/websocket_policy_client.py @@ -0,0 +1,196 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); +# Implemented by [Jinhui YE / HKUST University] in [2025]. +# +# Vendored from starVLA: deployment/model_server/tools/websocket_policy_client.py (2026-08-19). +# Source repo: https://github.com/starVLA/starVLA + +import logging +import os +import time +from typing import Any, Dict, List, Optional, Tuple + +import websockets.sync.client +from typing_extensions import override + +from . import msgpack_numpy + +# ============================================================================= +# TRAIN / TEST CONSISTENCY REMINDER (shown at every eval entry point) +# ----------------------------------------------------------------------------- +# Every eval benchmark under `examples/` connects to the policy server through +# this client, so this banner is emitted once per eval run. Embodied policies +# are extremely sensitive to the gap between how observations are built during +# TRAINING versus INFERENCE. A silent mismatch will NOT raise an error, it will +# only quietly degrade the success rate. +# ============================================================================= +_CONSISTENCY_REMINDER = ( + "\n" + "============================================================\n" + " [TRAIN/TEST CONSISTENCY CHECK] read before trusting results\n" + "------------------------------------------------------------\n" + " Make sure the EVAL observation matches TRAINING for:\n" + " - state : whether proprioceptive state is fed (use_state)\n" + " and its dimension / ordering\n" + " - img size : resize / crop resolution (e.g. 224x224)\n" + " - img count: how many camera views are fed to the model\n" + " - img order: the ordering of those camera views\n" + " - horizon : action chunk size / action horizon\n" + " A mismatch on ANY of these silently lowers the success rate.\n" + " Cross-check the values below against your training config.\n" + " The client will NOT infer or reorder camera views for you.\n" + "============================================================" +) + + +def _as_image_sequence(value: Any) -> Optional[List[Any]]: + if value is None: + return None + if isinstance(value, (list, tuple)): + return list(value) + return [value] + + +def _image_hw(image: Any) -> Optional[Tuple[int, int]]: + shape = getattr(image, "shape", None) + if shape is not None and len(shape) >= 2: + return int(shape[0]), int(shape[1]) + size = getattr(image, "size", None) + if isinstance(size, tuple) and len(size) >= 2: + return int(size[1]), int(size[0]) + return None + + +def _expected_image_hw(metadata: Dict) -> Optional[Tuple[int, int]]: + size = metadata.get("training_obs_image_size") + if isinstance(size, (list, tuple)) and len(size) == 2: + return int(size[0]), int(size[1]) + return None + + +class WebsocketClientPolicy: + """Implements the Policy interface by communicating with a server over websocket. + + See WebsocketPolicyServer for a corresponding server implementation. + """ + + def __init__(self, host: str = "127.0.0.1", port: Optional[int] = 10093, api_key: Optional[str] = None) -> None: + # 0.0.0.0 cannot be used as a connection target, here default 127.0.0.1 + self._uri = f"ws://{host}" + if port is not None: + self._uri += f":{port}" + self._packer = msgpack_numpy.Packer() + self._api_key = api_key + self._ws, self._server_metadata = self._wait_for_server() + self._did_log_eval_observation_contract = False + self._did_warn_eval_observation_mismatch = False + + # Remind the user to keep the eval-time observation pipeline aligned with + # training, and echo the server metadata so the values can be verified. + logging.warning(_CONSISTENCY_REMINDER) + logging.warning("[TRAIN/TEST CONSISTENCY CHECK] server metadata: %s", self._server_metadata) + + def get_server_metadata(self) -> Dict: + return self._server_metadata + + def _wait_for_server(self, timeout: float = 300) -> Tuple[websockets.sync.client.ClientConnection, Dict]: + logging.info(f"Waiting for server at {self._uri}...") + start_time = time.time() + + for k in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy"): + os.environ.pop(k, None) + + while True: + if time.time() - start_time > timeout: + raise TimeoutError(f"Failed to connect to server within {timeout} seconds") + + try: + headers = {"Authorization": f"Api-Key {self._api_key}"} if self._api_key else None + conn = websockets.sync.client.connect( + self._uri, + compression=None, + max_size=None, + additional_headers=headers, + open_timeout=150, + ping_interval=None, + ping_timeout=60, + ) + metadata = msgpack_numpy.unpackb(conn.recv()) + return conn, metadata + except ConnectionRefusedError: + logging.info(f"Still waiting for server {self._uri} ...") + time.sleep(2) + + def close(self) -> None: + try: + self._ws.close() + except Exception: + pass + + @override + def predict_action(self, query_info: Dict) -> Dict: + self._check_eval_observation_contract(query_info) + data = self._packer.pack(query_info) + self._ws.send(data) + response = self._ws.recv() + if isinstance(response, str): + raise RuntimeError(f"Error in inference server:\n{response}") + return msgpack_numpy.unpackb(response) + + def _check_eval_observation_contract(self, query_info: Dict) -> None: + examples = query_info.get("examples") + if not isinstance(examples, list) or not examples: + return + + expected_hw = _expected_image_hw(self._server_metadata) + image_counts: List[int] = [] + image_shapes: List[List[Optional[Tuple[int, int]]]] = [] + + for idx, example in enumerate(examples): + if not isinstance(example, dict): + continue + images = _as_image_sequence(example.get("image")) + if images is None: + logging.warning( + "[TRAIN/TEST CONSISTENCY CHECK] example %d has no `image` key. " + "Verify this is intended for the checkpoint metadata=%s", + idx, + self._server_metadata, + ) + continue + shapes = [_image_hw(img) for img in images] + image_counts.append(len(images)) + image_shapes.append(shapes) + + if expected_hw is not None: + for image_idx, hw in enumerate(shapes): + if hw is not None and hw != expected_hw: + self._warn_eval_observation_mismatch_once( + "[TRAIN/TEST CONSISTENCY CHECK] eval image size mismatch: " + f"example={idx}, image_index={image_idx}, got={hw}, " + f"training_obs_image_size={expected_hw}, metadata={self._server_metadata}. " + "Resize/crop explicitly in the benchmark interface before calling predict_action." + ) + + if image_counts and len(set(image_counts)) > 1: + self._warn_eval_observation_mismatch_once( + "[TRAIN/TEST CONSISTENCY CHECK] inconsistent image counts across eval batch: " + f"image_counts={image_counts}. Each example should use the same explicit camera contract." + ) + + if not self._did_log_eval_observation_contract and image_counts: + self._did_log_eval_observation_contract = True + logging.info( + "[TRAIN/TEST CONSISTENCY CHECK] eval request image_count=%s image_shapes=%s " + "server_metadata=%s. The benchmark interface is responsible for camera order; " + "verify this order manually against the checkpoint training setup.", + image_counts[0], + image_shapes[0] if image_shapes else None, + self._server_metadata, + ) + + def _warn_eval_observation_mismatch_once(self, message: str) -> None: + if self._did_warn_eval_observation_mismatch: + return + self._did_warn_eval_observation_mismatch = True + logging.warning(message)