xctrl server: bread convention is 20D state in / 14D actions out (tofu template was 14D state)

deploy/serve_bread_xctrl.py

@@ -1,16 +1,21 @@
 """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.
+Modeled on Tianming's poke/serve_xctrl_openpi_policy.py -- BUT that template
+is from the cooking (tofu) project whose state is 14D. The bread convention
+agreed with Zhengmao is **20D state in / 14D actions out**, which matches
+this model exactly, so the shim below is mostly key renaming.
 
 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
+        "state": ndarray(20,),  # per arm [xyz(3), rot6d(6), grip(1)], L then R,
+                                # grasp_site frame.
+                                # DEFAULT ASSUMPTION: already relative_to_first
+                                # (Zhengmao's yam_umi convention).      <-- CONFIRM
         "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
+        "prompt": str,          # optional; defaults to the training prompt
+        "reset": bool,          # optional; only used in --state-mode=absolute
     }
 
 Server returns:
@@ -21,19 +26,16 @@ Server returns:
         ... (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.
+--state-mode:
+    relative  (default) executor already re-expresses each arm's pose in its
+              episode-start frame -> passthrough.
+    absolute  executor sends absolute poses -> this server records the first
+              pose after start/{"reset": True} and re-expresses (STATEFUL:
+              send reset at every new episode or actions will be garbage).
 
 Usage:
     uv run python deploy/serve_bread_xctrl.py \
-        --port=8000 \
+        --port=8000 --state-mode=relative \
         policy:checkpoint \
         --policy.config=pi05_bread_r1_lora \
         --policy.dir=/PATH/TO/pi05-bread-r1-top200/10000
@@ -44,11 +46,10 @@ from __future__ import annotations
 import dataclasses
 import logging
 import socket
-from typing import Any
+from typing import Any, Literal
 
 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
@@ -57,10 +58,17 @@ 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."""
+def _rot6d_to_matrix(r6: np.ndarray) -> np.ndarray:
+    a, b = r6[:3], r6[3:]
+    x = a / np.linalg.norm(a)
+    b = b - np.dot(x, b) * x
+    y = b / np.linalg.norm(b)
+    return np.stack([x, y, np.cross(x, y)], axis=1)
+
+
+def _pose_from_10d(x: np.ndarray) -> np.ndarray:
     T = np.eye(4)
-    T[:3, :3] = Rotation.from_rotvec(x[3:6]).as_matrix()
+    T[:3, :3] = _rot6d_to_matrix(x[3:9])
     T[:3, 3] = x[:3]
     return T
 
@@ -79,31 +87,35 @@ class Checkpoint:
 class Args:
     port: int = 8000
     default_prompt: str = TRAIN_PROMPT
+    state_mode: Literal["relative", "absolute"] = "relative"
     policy: Checkpoint = dataclasses.field(default_factory=Checkpoint)
 
 
 class BreadPolicyServer:
-    """Flat xctrl obs -> BreadInputs keys, with stateful relative_to_first."""
+    """Flat xctrl obs -> BreadInputs keys (20D state passthrough or re-expression)."""
 
-    def __init__(self, openpi_policy, default_prompt: str = TRAIN_PROMPT) -> None:
+    def __init__(self, openpi_policy, state_mode: str = "relative",
+                 default_prompt: str = TRAIN_PROMPT) -> None:
         self._policy = openpi_policy
+        self._state_mode = state_mode
         self._default_prompt = default_prompt
-        self._start_inv: list[np.ndarray] | None = None  # [inv(T_L0), inv(T_R0)]
+        self._start_inv: list[np.ndarray] | None = None  # absolute mode only
 
-    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]]
+    def _state(self, state20: np.ndarray, reset: bool) -> np.ndarray:
+        state20 = np.asarray(state20, dtype=np.float64).reshape(20)
+        if self._state_mode == "relative":
+            return state20.astype(np.float32)
+        poses = [_pose_from_10d(state20[0:10]), _pose_from_10d(state20[10:20])]
         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)):
+        out = state20.astype(np.float32).copy()
+        for i, (inv0, T) in enumerate(zip(self._start_inv, poses)):
             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
+            # gripper (off+9) passes through unchanged
         return out
 
     def infer(self, obs: dict[str, Any]) -> dict[str, Any]:
@@ -111,7 +123,7 @@ class BreadPolicyServer:
         if isinstance(prompt, bytes):
             prompt = prompt.decode("utf-8")
         openpi_obs = {
-            "state": self._state_20d(obs["state"], bool(obs.get("reset", False))),
+            "state": self._state(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),
@@ -134,10 +146,11 @@ def main(args: Args) -> None:
         _policy_config.create_trained_policy(
             _config.get_config(args.policy.config), args.policy.dir,
             default_prompt=args.default_prompt),
+        state_mode=args.state_mode,
         default_prompt=args.default_prompt)
     hostname = socket.gethostname()
-    logging.info("bread policy server on %s (%s):%d", hostname,
-                 socket.gethostbyname(hostname), args.port)
+    logging.info("bread policy server on %s (%s):%d state_mode=%s", hostname,
+                 socket.gethostbyname(hostname), args.port, args.state_mode)
     websocket_policy_server.WebsocketPolicyServer(
         policy=policy, host="0.0.0.0", port=args.port).serve_forever()