moonspace release v1

Author: hlink

example/deploy/panthera_single_client.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+"""Run Panthera single-arm real-robot inference through ModelClient.
+
+The TCP model server side should expose an RPC method such as ``infer(obs)``
+that returns either a single action [14] or an action chunk [T, 14]. For this
+single-arm robot, the first 7 dims are used as left_arm joint(6)+gripper(1),
+and the remaining right-arm dims are ignored.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+import time
+from typing import Any
+
+import numpy as np
+
+THIS_FILE = Path(__file__).resolve()
+DEPLOY_DIR = THIS_FILE.parent
+REPO_ROOT = DEPLOY_DIR.parent.parent
+
+for path in (DEPLOY_DIR, REPO_ROOT, REPO_ROOT / "src"):
+    path_str = str(path)
+    if path_str not in sys.path:
+        sys.path.insert(0, path_str)
+
+from robot_client_server.model_client import ModelClient
+from my_robot.panthera_single import PantheraSingle
+
+
+ACT_STATE_DIM = 14
+LEFT_ACTION_DIM = 7
+
+
+def _color_to_chw(color: Any) -> np.ndarray:
+    image = np.asarray(color)
+    if image.ndim != 3 or image.shape[2] != 3:
+        raise ValueError(f"expected HWC RGB/BGR image, got shape={image.shape}")
+    return np.transpose(image.astype(np.uint8, copy=False), (2, 0, 1))
+
+
+def _build_state(controller_data: dict[str, Any]) -> np.ndarray:
+    left = controller_data["left_arm"]
+    joint = np.asarray(left["joint"], dtype=np.float32).reshape(-1)
+    if joint.shape != (6,):
+        raise ValueError(f"left_arm.joint must be shape (6,), got {joint.shape}")
+    gripper = np.asarray(left["gripper"], dtype=np.float32).reshape(1)
+    state = np.zeros((ACT_STATE_DIM,), dtype=np.float32)
+    state[:6] = joint
+    state[6] = gripper[0]
+    return state
+
+
+def build_observation(raw_data: Any, prompt: str) -> dict[str, Any]:
+    controller_data, sensor_data = raw_data
+    state = _build_state(controller_data)
+    cam_high = _color_to_chw(sensor_data["cam_head"]["color"])
+    cam_wrist = _color_to_chw(sensor_data["cam_wrist"]["color"])
+
+    # Keep cam_wrist for single-arm checkpoints. Also provide the dual-arm keys as
+    # aliases so older OpenPI/ALOHA wrappers that still look for wrist-left/right
+    # can consume the same observation without changing this robot client.
+    return {
+        "state": state,
+        "images": {
+            "cam_high": cam_high,
+            "cam_wrist": cam_wrist,
+            "cam_right_wrist": cam_wrist,
+            "cam_left_wrist": np.zeros_like(cam_wrist),
+        },
+        "prompt": prompt,
+    }
+
+
+def _first_action(action_result: Any) -> np.ndarray:
+    actions = np.asarray(action_result, dtype=np.float32)
+    if actions.ndim == 0:
+        raise ValueError("model returned scalar action")
+    if actions.ndim == 1:
+        action = actions
+    else:
+        action = actions[0]
+    if action.shape[0] < LEFT_ACTION_DIM:
+        raise ValueError(f"action must have at least {LEFT_ACTION_DIM} dims, got {action.shape}")
+    return action
+
+
+def action_to_move_data(action: np.ndarray, *, gripper_min: float, gripper_max: float) -> dict[str, Any]:
+    left = np.asarray(action[:LEFT_ACTION_DIM], dtype=np.float32).reshape(-1)
+    joint = left[:6]
+    gripper = float(np.clip(left[6], gripper_min, gripper_max))
+    return {
+        "arm": {
+            "left_arm": {
+                "joint": joint,
+                "gripper": gripper,
+            }
+        }
+    }
+
+
+def maybe_call(client: ModelClient, func_name: str) -> None:
+    try:
+        client.call(func_name=func_name)
+    except Exception as exc:
+        print(f"[panthera-client] ignore {func_name!r} RPC failure: {exc}")
+
+
+def run(args: argparse.Namespace) -> None:
+    robot = PantheraSingle(enable_cameras=not args.no_cameras)
+    client = ModelClient(host=args.host, port=args.port, timeout=args.timeout)
+    period = 1.0 / args.hz
+
+    try:
+        robot.set_up()
+        if not args.no_reset:
+            robot.reset()
+        maybe_call(client, args.reset_rpc)
+
+        print("[panthera-client] ready. Press ENTER to start inference, Ctrl+C to stop.")
+        input()
+        print("[panthera-client] running...")
+
+        step = 0
+        while args.max_steps <= 0 or step < args.max_steps:
+            loop_start = time.monotonic()
+            raw = robot.get()
+            obs = build_observation(raw, args.prompt)
+            action_result = client.call(func_name=args.infer_rpc, obs=obs)
+            action = _first_action(action_result)
+            move_data = action_to_move_data(
+                action,
+                gripper_min=args.gripper_min,
+                gripper_max=args.gripper_max,
+            )
+
+            if args.dry_run:
+                print(f"[panthera-client] step={step} action={np.array2string(action[:7], precision=4)}")
+            else:
+                robot.move(move_data)
+
+            step += 1
+            if args.log_every > 0 and step % args.log_every == 0:
+                print(f"[panthera-client] step={step}")
+
+            elapsed = time.monotonic() - loop_start
+            sleep_time = period - elapsed
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+
+    except KeyboardInterrupt:
+        print("\n[panthera-client] interrupted by user")
+    finally:
+        try:
+            robot.shutdown(damped=args.brake_on_exit)
+        finally:
+            client.close()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Panthera single-arm model-client deployment loop.")
+    parser.add_argument("--host", default="localhost", help="model server host")
+    parser.add_argument("--port", type=int, default=9999, help="model server TCP port")
+    parser.add_argument("--timeout", type=float, default=30.0, help="socket timeout seconds")
+    parser.add_argument("--infer-rpc", default="infer", help="model server inference method name")
+    parser.add_argument("--reset-rpc", default="reset_episode", help="model server episode reset method name")
+    parser.add_argument("--prompt", default="", help="language prompt sent with each observation")
+    parser.add_argument("--hz", type=float, default=10.0, help="control loop frequency")
+    parser.add_argument("--max-steps", type=int, default=0, help="0 means run until Ctrl+C")
+    parser.add_argument("--log-every", type=int, default=10, help="print progress every N steps; <=0 disables")
+    parser.add_argument("--gripper-min", type=float, default=0.0, help="minimum normalized gripper command")
+    parser.add_argument("--gripper-max", type=float, default=1.0, help="maximum normalized gripper command")
+    parser.add_argument("--dry-run", action="store_true", help="run inference but do not move the robot")
+    parser.add_argument("--no-reset", action="store_true", help="do not reset robot before starting")
+    parser.add_argument("--no-cameras", action="store_true", help="initialize robot without cameras")
+    parser.add_argument("--brake-on-exit", action="store_true", help="brake instead of stop on exit")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    run(parse_args())