Add xctrl-style flat-obs policy server (Tianming template)

deploy/serve_bread_xctrl.py

@@ -0,0 +1,147 @@
+"""Serve the bread R1 policy (pi05_bread_r1_lora) over WebSocket, xctrl-style.
+
+Modeled on Tianming's poke/serve_xctrl_openpi_policy.py flat-obs contract.
+
+Client sends a flat obs dict:
+    {
+        "state": ndarray(14,),   # per arm [xyz(3), axis_angle(3), grip(1)], L then R
+                                 # ABSOLUTE grasp_site poses            <-- CONFIRM layout/frame
+        "camera_ego": ndarray(H, W, 3),   # accepted but IGNORED (model's base slot is masked)
+        "camera_left": ndarray(H, W, 3),
+        "camera_right": ndarray(H, W, 3),
+        "prompt": str,           # optional; defaults to the training prompt
+        "reset": bool,           # optional; True on the first frame of an episode
+    }
+
+Server returns:
+    {
+        "actions": ndarray(40, 14),  # per arm [dxyz(3), axis-angle drot(3), ABS grip(1)],
+                                     # stepwise-chained deltas @30Hz:
+                                     # T_cmd,k = T_cmd,k-1 @ delta_k from T_current
+        ... (openpi policy timing fields passed through)
+    }
+
+Differences vs the cooking server -- all handled inside this shim:
+  1. key names: our BreadInputs takes left_wrist_image / right_wrist_image /
+     state / prompt (no ego/base camera -- it was a human egocam in training).
+  2. state conversion: client 14D axis-angle -> model 20D rot6d.
+  3. relative_to_first: the model was trained on states re-expressed in each
+     arm's EPISODE-START frame. This server is therefore STATEFUL: it records
+     the first absolute pose after start/reset and re-expresses every
+     subsequent state. Send {"reset": True} at every new episode, or restart
+     the server. Getting this wrong produces confidently wrong actions.
+
+Usage:
+    uv run python deploy/serve_bread_xctrl.py \
+        --port=8000 \
+        policy:checkpoint \
+        --policy.config=pi05_bread_r1_lora \
+        --policy.dir=/PATH/TO/pi05-bread-r1-top200/10000
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import socket
+from typing import Any
+
+import numpy as np
+import tyro
+from scipy.spatial.transform import Rotation
+
+from openpi.policies import policy_config as _policy_config
+from openpi.serving import websocket_policy_server  # swap for xctrl.policies.WebSocketPolicyServer inside xctrl
+from openpi.training import config as _config
+
+TRAIN_PROMPT = "take the bread out of the bowl and put the bread into the toaster"
+
+
+def _pose_from_7d(x: np.ndarray) -> np.ndarray:
+    """[xyz(3), axis_angle(3)] -> 4x4."""
+    T = np.eye(4)
+    T[:3, :3] = Rotation.from_rotvec(x[3:6]).as_matrix()
+    T[:3, 3] = x[:3]
+    return T
+
+
+def _rot6d(R: np.ndarray) -> np.ndarray:
+    return np.concatenate([R[:, 0], R[:, 1]])
+
+
+@dataclasses.dataclass
+class Checkpoint:
+    config: str = "pi05_bread_r1_lora"
+    dir: str = ""
+
+
+@dataclasses.dataclass
+class Args:
+    port: int = 8000
+    default_prompt: str = TRAIN_PROMPT
+    policy: Checkpoint = dataclasses.field(default_factory=Checkpoint)
+
+
+class BreadPolicyServer:
+    """Flat xctrl obs -> BreadInputs keys, with stateful relative_to_first."""
+
+    def __init__(self, openpi_policy, default_prompt: str = TRAIN_PROMPT) -> None:
+        self._policy = openpi_policy
+        self._default_prompt = default_prompt
+        self._start_inv: list[np.ndarray] | None = None  # [inv(T_L0), inv(T_R0)]
+
+    def _state_20d(self, state14: np.ndarray, reset: bool) -> np.ndarray:
+        state14 = np.asarray(state14, dtype=np.float64).reshape(14)
+        poses = [_pose_from_7d(state14[0:6]), _pose_from_7d(state14[7:13])]
+        grips = [state14[6], state14[13]]
+        if reset or self._start_inv is None:
+            self._start_inv = [np.linalg.inv(p) for p in poses]
+            logging.info("relative_to_first anchor set (episode start)")
+        out = np.zeros(20, dtype=np.float32)
+        for i, (inv0, T, g) in enumerate(zip(self._start_inv, poses, grips)):
+            rel = inv0 @ T
+            off = i * 10
+            out[off:off + 3] = rel[:3, 3]
+            out[off + 3:off + 9] = _rot6d(rel[:3, :3])
+            out[off + 9] = g
+        return out
+
+    def infer(self, obs: dict[str, Any]) -> dict[str, Any]:
+        prompt = obs.get("prompt", self._default_prompt)
+        if isinstance(prompt, bytes):
+            prompt = prompt.decode("utf-8")
+        openpi_obs = {
+            "state": self._state_20d(obs["state"], bool(obs.get("reset", False))),
+            "left_wrist_image": np.asarray(obs["camera_left"]),
+            "right_wrist_image": np.asarray(obs["camera_right"]),
+            "prompt": str(prompt),
+            # camera_ego intentionally dropped: the model's base_0_rgb slot is
+            # zeroed+masked (training top view was a human egocam).
+        }
+        result = self._policy.infer(openpi_obs)
+        result["actions"] = np.asarray(result["actions"])[:, :14]
+        return result
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        return getattr(self._policy, "metadata", {}) or {}
+
+
+def main(args: Args) -> None:
+    if not args.policy.dir:
+        raise ValueError("policy.dir is required (…/pi05-bread-r1-top200/10000)")
+    policy = BreadPolicyServer(
+        _policy_config.create_trained_policy(
+            _config.get_config(args.policy.config), args.policy.dir,
+            default_prompt=args.default_prompt),
+        default_prompt=args.default_prompt)
+    hostname = socket.gethostname()
+    logging.info("bread policy server on %s (%s):%d", hostname,
+                 socket.gethostbyname(hostname), args.port)
+    websocket_policy_server.WebsocketPolicyServer(
+        policy=policy, host="0.0.0.0", port=args.port).serve_forever()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    main(tyro.cli(Args))