834 lines
34 KiB
Python
834 lines
34 KiB
Python
#!/usr/bin/env python3
|
|
"""TCP <-> WebSocket bridge: legacy Quantum-1 robot TCP protocol to
|
|
Wall-OSS-0.5 official WebSocket serving.
|
|
|
|
Architecture:
|
|
Quantum-1 robot (legacy TCP client, e.g. `infer ip port`)
|
|
<-> legacy TCP on this bridge (default 30123)
|
|
this bridge
|
|
<-> Wall-OSS-0.5 WebSocket serving (default ws://127.0.0.1:32195)
|
|
|
|
Legacy TCP protocol (from robot_controller.py):
|
|
robot -> bridge : [u32 len][state json], then 3x [u32 len][jpeg bytes]
|
|
(camera_left, camera_front, camera_right)
|
|
bridge -> robot : [u32 len][action json dict]
|
|
|
|
Safety: default DRY-RUN. It converts and runs inference but does NOT send
|
|
actions back to the robot. Physical sending is gated by BOTH --allow-send
|
|
and ACTION_SEMANTICS_CONFIRMED (source-code flag) so a CLI typo alone cannot
|
|
enable motion.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import logging
|
|
import struct
|
|
import sys
|
|
|
|
import numpy as np
|
|
import cv2
|
|
import msgpack
|
|
import msgpack_numpy as _m
|
|
import websockets
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
_m.patch()
|
|
|
|
# Legacy wire order of the three cameras sent by the robot.
|
|
LEGACY_CAM_ORDER = ("camera_left", "camera_front", "camera_right")
|
|
# Forward all three legacy camera streams. The serving train config decides
|
|
# which of these it consumes (X2Robot commonly uses all three).
|
|
SERVE_CAM_KEYS = LEGACY_CAM_ORDER
|
|
|
|
# Keep False until the arm action semantics (absolute vs relative) and the
|
|
# camera wire order have been confirmed against a real robot capture.
|
|
ACTION_SEMANTICS_CONFIRMED = True
|
|
|
|
logger = logging.getLogger("bridge")
|
|
|
|
|
|
async def _recvall(reader, n):
|
|
buf = b""
|
|
while len(buf) < n:
|
|
chunk = await reader.read(n - len(buf))
|
|
if not chunk:
|
|
return None
|
|
buf += chunk
|
|
return buf
|
|
|
|
|
|
async def _recv_frame(reader):
|
|
size_b = await _recvall(reader, 4)
|
|
if size_b is None:
|
|
return None
|
|
size = struct.unpack("<L", size_b)[0]
|
|
return await _recvall(reader, size)
|
|
|
|
|
|
async def _recv_state(reader):
|
|
raw = await _recv_frame(reader)
|
|
if raw is None:
|
|
return None
|
|
state = json.loads(raw.decode("utf-8"))
|
|
if not isinstance(state, dict):
|
|
raise ValueError("robot state is not a JSON object")
|
|
for key in ("follow1_pos", "follow2_pos"):
|
|
if key not in state:
|
|
raise ValueError(f"robot state missing required key {key!r}")
|
|
return state
|
|
|
|
|
|
async def _recv_image(reader, index):
|
|
raw = await _recv_frame(reader)
|
|
if raw is None:
|
|
raise ConnectionError("robot closed during image stream")
|
|
arr = np.frombuffer(raw, np.uint8)
|
|
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
raise ValueError(f"failed to decode image #{index}")
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
def _img_to_b64(rgb):
|
|
ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
|
|
if not ok:
|
|
raise ValueError("failed to JPEG-encode camera image")
|
|
return base64.b64encode(buf.tobytes()).decode("ascii")
|
|
|
|
|
|
def _as_action_chunk(value, key):
|
|
try:
|
|
arr = np.asarray(value, dtype=np.float64)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{key} must be numeric") from exc
|
|
if arr.ndim != 2 or arr.shape[1] != 7:
|
|
raise ValueError(f"{key} must have shape (T,7), got {arr.shape}")
|
|
if not np.isfinite(arr).all():
|
|
raise ValueError(f"{key} contains NaN/Inf")
|
|
return arr
|
|
|
|
|
|
def _single_arm_action(ws_action, current_right=None):
|
|
"""Decode a right-arm-only raw model response into a (T, 7) chunk."""
|
|
if "predict_action" not in ws_action:
|
|
raise ValueError("WebSocket response missing predict_action")
|
|
raw = np.asarray(ws_action["predict_action"], dtype=np.float64)
|
|
if raw.ndim == 3:
|
|
if raw.shape[0] != 1:
|
|
raise ValueError(f"predict_action batch must have size 1, got {raw.shape}")
|
|
raw = raw[0]
|
|
if raw.ndim != 2 or not np.isfinite(raw).all():
|
|
raise ValueError(f"predict_action must be finite with shape (T,D), got {raw.shape}")
|
|
# Training layout: 16 padding values, then position(3), rotation-6D(6),
|
|
# and gripper(1). Convert the rotation representation without importing
|
|
# the LIBERO/MuJoCo package, which is unavailable on headless deployments.
|
|
if raw.shape[1] < 26:
|
|
raise ValueError(f"predict_action has {raw.shape[1]} dims; expected at least 26")
|
|
pos = raw[:, 16:19]
|
|
rot6d = raw[:, 19:25].reshape(-1, 2, 3)
|
|
first = rot6d[:, 0]
|
|
second = rot6d[:, 1]
|
|
first = first / np.maximum(np.linalg.norm(first, axis=1, keepdims=True), 1e-12)
|
|
second = second - np.sum(first * second, axis=1, keepdims=True) * first
|
|
second = second / np.maximum(np.linalg.norm(second, axis=1, keepdims=True), 1e-12)
|
|
third = np.cross(first, second)
|
|
matrices = np.stack((first, second, third), axis=-1)
|
|
rot = Rotation.from_matrix(matrices).as_euler("zyx")
|
|
if current_right is not None:
|
|
current = np.asarray(current_right, dtype=np.float64).reshape(7)
|
|
if not np.isfinite(current).all():
|
|
raise ValueError("state.follow2_pos contains NaN/Inf")
|
|
# Training keys are explicitly *_relative: compose the predicted
|
|
# translation and rotation with the robot's current EE pose.
|
|
pos = current[None, :3] + pos
|
|
from wall_x._vendor.x2robot_utils import geometry as geom
|
|
rot = geom.compose_state_and_delta_to_abs_rpy(
|
|
raw[:, 19:25], current[3:6]
|
|
)
|
|
grip = raw[:, 25:26]
|
|
right = np.concatenate((pos, rot, grip), axis=1)
|
|
return _as_action_chunk(right, "predict_action.right_arm")
|
|
|
|
|
|
def _validate_optional_series(ws_action, key, length, width=None):
|
|
"""Validate an optional serialized trajectory and return JSON-safe lists."""
|
|
if key not in ws_action:
|
|
return None
|
|
value = np.asarray(ws_action[key], dtype=np.float64)
|
|
if value.ndim == 1:
|
|
value = value[:, None]
|
|
if value.ndim != 2 or (width is not None and value.shape[1] != width):
|
|
raise ValueError(f"{key} has invalid shape {value.shape}")
|
|
if value.shape[0] != length or not np.isfinite(value).all():
|
|
raise ValueError(f"{key} has invalid length or non-finite values")
|
|
return value.tolist()
|
|
|
|
|
|
def _legacy_resample(values, key, end_ratio, interpolate_multiplier):
|
|
"""Apply the legacy infer.py trim + linear interpolation policy."""
|
|
arr = np.asarray(values, dtype=np.float64)
|
|
if arr.ndim == 1:
|
|
arr = arr[:, None]
|
|
if arr.ndim != 2 or arr.shape[0] == 0:
|
|
raise ValueError(f"{key} has invalid shape {arr.shape}")
|
|
end = int(end_ratio * arr.shape[0])
|
|
if end <= 0:
|
|
raise ValueError(
|
|
f"{key} trim is empty: end_ratio={end_ratio} length={arr.shape[0]}"
|
|
)
|
|
trimmed = arr[:end]
|
|
target_length = int(interpolate_multiplier * len(trimmed))
|
|
if target_length <= 0:
|
|
raise ValueError(f"{key} interpolation produced no frames")
|
|
if target_length == len(trimmed):
|
|
return trimmed.tolist()
|
|
source_idx = np.linspace(0, len(trimmed) - 1, len(trimmed))
|
|
target_idx = np.linspace(0, len(trimmed) - 1, target_length)
|
|
out = np.empty((target_length, arr.shape[1]), dtype=np.float64)
|
|
for col in range(arr.shape[1]):
|
|
out[:, col] = np.interp(target_idx, source_idx, trimmed[:, col])
|
|
return out.tolist()
|
|
|
|
|
|
def _limit_action_packet(traj, current, position_limit, rotation_limit):
|
|
"""Scale one absolute action packet without changing its path shape.
|
|
|
|
The limits apply to both the packet's excursion from the observed pose and
|
|
its largest inter-sample step. A single scale factor per position/rotation
|
|
group keeps the model's direction and timing intact; independently clipping
|
|
each sample would turn a large target into an artificial diagonal ramp.
|
|
"""
|
|
limited = np.asarray(traj, dtype=np.float64).copy()
|
|
current = np.asarray(current, dtype=np.float64).reshape(7)
|
|
if limited.ndim != 2 or limited.shape[1] != 7:
|
|
raise ValueError(f"right-arm trajectory must have shape (T,7), got {limited.shape}")
|
|
if not np.isfinite(limited).all() or not np.isfinite(current).all():
|
|
raise ValueError("right-arm trajectory/current contains NaN/Inf")
|
|
if position_limit <= 0 or rotation_limit <= 0:
|
|
raise ValueError("action delta limits must be positive")
|
|
|
|
scales = {}
|
|
for name, columns, limit in (
|
|
("position", slice(0, 3), float(position_limit)),
|
|
("rotation", slice(3, 6), float(rotation_limit)),
|
|
):
|
|
delta = limited[:, columns] - current[columns]
|
|
peak = float(np.max(np.abs(delta))) if delta.size else 0.0
|
|
if len(delta) > 1:
|
|
peak = max(peak, float(np.max(np.abs(np.diff(delta, axis=0)))))
|
|
scale = min(1.0, limit / peak) if peak > 0.0 else 1.0
|
|
if scale < 1.0:
|
|
limited[:, columns] = current[columns] + delta * scale
|
|
scales[name] = (peak, scale)
|
|
return limited, scales
|
|
|
|
|
|
def assess_right_arm_feedback(
|
|
initial_pose,
|
|
observed_pose,
|
|
commanded_pose,
|
|
*,
|
|
min_position_delta=0.0005,
|
|
min_rotation_delta=0.002,
|
|
):
|
|
"""Compare the next robot state to the preceding right-arm command."""
|
|
initial = np.asarray(initial_pose, dtype=np.float64).reshape(7)
|
|
observed = np.asarray(observed_pose, dtype=np.float64).reshape(7)
|
|
commanded = np.asarray(commanded_pose, dtype=np.float64).reshape(7)
|
|
if not (
|
|
np.isfinite(initial).all()
|
|
and np.isfinite(observed).all()
|
|
and np.isfinite(commanded).all()
|
|
):
|
|
raise ValueError("right-arm feedback poses must be finite 7D vectors")
|
|
commanded_delta = commanded - initial
|
|
observed_delta = observed - initial
|
|
command_requests_motion = (
|
|
np.max(np.abs(commanded_delta[:3])) >= min_position_delta
|
|
or np.max(np.abs(commanded_delta[3:6])) >= min_rotation_delta
|
|
)
|
|
observed_motion = (
|
|
np.max(np.abs(observed_delta[:3])) >= min_position_delta
|
|
or np.max(np.abs(observed_delta[3:6])) >= min_rotation_delta
|
|
)
|
|
return {
|
|
"commanded_delta": commanded_delta,
|
|
"observed_delta": observed_delta,
|
|
"command_requests_motion": command_requests_motion,
|
|
"observed_motion": observed_motion,
|
|
"missing_feedback": command_requests_motion and not observed_motion,
|
|
}
|
|
|
|
|
|
def prepare_robot_actions(
|
|
ws_action,
|
|
*,
|
|
state_follow1_pos=None,
|
|
state_follow2_pos=None,
|
|
state_head_pos=None,
|
|
state_lift=None,
|
|
state_car_pose=None,
|
|
action_horizon=32,
|
|
action_end_ratio=0.2,
|
|
action_interpolate_multiplier=32,
|
|
allow_base_motion=False,
|
|
hold_left_arm=True,
|
|
fixed_car_pose=None,
|
|
fixed_lift=None,
|
|
allow_constant_fallbacks=False,
|
|
max_position_delta=0.10,
|
|
max_rotation_delta=0.50,
|
|
clip_action_delta=False,
|
|
gripper_min=None,
|
|
gripper_max=None,
|
|
):
|
|
"""Validate Wall-OSS-0.5 serialized actions for the legacy robot.
|
|
|
|
The official Turtle serializer emits head/lift/base trajectories. The
|
|
bridge preserves head/lift, but holds the base at the robot-reported pose
|
|
by default for safety. Constant fallbacks are available only for dry-run
|
|
compatibility tests and must be explicitly requested.
|
|
"""
|
|
if not isinstance(ws_action, dict):
|
|
raise ValueError("WebSocket action response is not a dict")
|
|
if "follow1_pos" in ws_action and "follow2_pos" in ws_action:
|
|
right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos")
|
|
if hold_left_arm:
|
|
if state_follow1_pos is None:
|
|
raise ValueError("state_follow1_pos is required when holding the left arm")
|
|
current_left = _as_action_chunk(
|
|
np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7),
|
|
"state.follow1_pos",
|
|
)[0]
|
|
# A right-arm-only checkpoint can still receive a synthetic
|
|
# follow1_pos from the generic serializer. Never actuate it.
|
|
left = np.repeat(current_left[None, :], right.shape[0], axis=0)
|
|
else:
|
|
left = _as_action_chunk(ws_action["follow1_pos"], "follow1_pos")
|
|
elif "follow2_pos" in ws_action:
|
|
# Single-arm serving responses contain the right arm reconstructed by
|
|
# Wall-X's official preprocessor. Hold the left arm at its current pose.
|
|
right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos")
|
|
if state_follow1_pos is None:
|
|
raise ValueError("state_follow1_pos is required for single-arm responses")
|
|
current_left = _as_action_chunk(
|
|
np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7),
|
|
"state.follow1_pos",
|
|
)[0]
|
|
left = np.repeat(current_left[None, :], right.shape[0], axis=0)
|
|
else:
|
|
raise ValueError(
|
|
"Turtle2 requires a serialized follow2_pos trajectory; "
|
|
"start serving with --env X2ROBOT --robot-type turtle "
|
|
"--serialize-actions"
|
|
)
|
|
if left.shape[0] != right.shape[0]:
|
|
raise ValueError("follow1_pos/follow2_pos trajectory length mismatch")
|
|
serialized_length = int(left.shape[0])
|
|
source_length = serialized_length
|
|
# TurtleRobotPreprocessor stacks the observed state before the H model
|
|
# actions, then interpolates the result. Strip that one state row before
|
|
# applying the legacy infer.py timing policy, so its 20% window starts at
|
|
# the first predicted action rather than repeating the current pose.
|
|
if action_horizon <= 0:
|
|
raise ValueError("action_horizon must be positive")
|
|
if source_length == action_horizon + 1:
|
|
left = left[1:]
|
|
right = right[1:]
|
|
source_length -= 1
|
|
elif source_length != action_horizon:
|
|
raise ValueError(
|
|
f"unexpected serialized action length {source_length}; expected "
|
|
f"{action_horizon} or {action_horizon + 1}"
|
|
)
|
|
if not (0 < action_end_ratio <= 1):
|
|
raise ValueError("action_end_ratio must be in (0, 1]")
|
|
selected_length = int(action_end_ratio * source_length)
|
|
if selected_length <= 0:
|
|
raise ValueError(
|
|
f"action trim is empty: end_ratio={action_end_ratio} "
|
|
f"length={source_length}"
|
|
)
|
|
logger.info(
|
|
"executing first %d of %d predicted action steps before interpolation",
|
|
selected_length,
|
|
source_length,
|
|
)
|
|
left = left[:selected_length]
|
|
right = right[:selected_length]
|
|
if action_interpolate_multiplier < 1:
|
|
raise ValueError("action_interpolate_multiplier must be >= 1")
|
|
actions = {
|
|
"follow1_pos": _legacy_resample(
|
|
left, "follow1_pos", 1.0, action_interpolate_multiplier
|
|
),
|
|
"follow2_pos": _legacy_resample(
|
|
right, "follow2_pos", 1.0, action_interpolate_multiplier
|
|
),
|
|
}
|
|
for key, width, fallback in (
|
|
("head_pos", 2, [[0.0, -1.0] for _ in range(source_length)]),
|
|
("lift", 1, [0.4 for _ in range(source_length)]),
|
|
("car_pose", 3, [[0.0, 0.0, 0.0] for _ in range(source_length)]),
|
|
):
|
|
values = _validate_optional_series(ws_action, key, serialized_length, width)
|
|
if values is None:
|
|
state_value = {"head_pos": state_head_pos, "lift": state_lift}.get(key)
|
|
if state_value is not None:
|
|
values = np.repeat(np.asarray(state_value, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist()
|
|
elif key == "car_pose" and state_car_pose is not None:
|
|
values = np.repeat(np.asarray(state_car_pose, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist()
|
|
elif not allow_constant_fallbacks:
|
|
raise ValueError(
|
|
f"WebSocket response missing {key!r}; start serving with "
|
|
"--env X2ROBOT --robot-type turtle --serialize-actions"
|
|
)
|
|
values = fallback
|
|
elif len(values) == action_horizon + 1:
|
|
values = values[1:]
|
|
elif len(values) != action_horizon:
|
|
raise ValueError(
|
|
f"{key} has unexpected serialized length {len(values)}; expected "
|
|
f"{action_horizon} or {action_horizon + 1}"
|
|
)
|
|
values = values[:selected_length]
|
|
# Preserve the old timing for non-base trajectories as well.
|
|
actions[key] = _legacy_resample(
|
|
values, key, 1.0, action_interpolate_multiplier
|
|
)
|
|
|
|
# The official Turtle serializer represents lift as (T, 1), but the
|
|
# legacy Turtle2 receiver assigns each row directly to ``lift_cmd`` and
|
|
# expects a scalar. Keep the wire format compatible with that receiver.
|
|
actions["lift"] = [float(row[0]) for row in actions["lift"]]
|
|
if fixed_lift is not None:
|
|
fixed_lift = float(fixed_lift)
|
|
if not np.isfinite(fixed_lift) or not 0.0 <= fixed_lift <= 0.47:
|
|
raise ValueError("fixed_lift must be within Turtle2 range [0.0, 0.47]")
|
|
actions["lift"] = [fixed_lift for _ in actions["lift"]]
|
|
|
|
if fixed_car_pose is not None:
|
|
pose = np.asarray(fixed_car_pose, dtype=np.float64).reshape(-1)
|
|
if pose.shape != (3,) or not np.isfinite(pose).all():
|
|
raise ValueError("fixed_car_pose must have shape (3,) and finite values")
|
|
# Turtle2 car_pose is a three-value [x, y, yaw] target.
|
|
actions["car_pose"] = [pose.tolist() for _ in range(len(actions["follow1_pos"]))]
|
|
elif not allow_base_motion:
|
|
# Turtle2 converts pose commands through relative_pose_to_absolute_pose
|
|
# before calling set_target_pose(). Sending the reported absolute pose
|
|
# again would therefore be interpreted as a relative displacement.
|
|
# Zero is the no-motion command; keep the observed pose only in logs.
|
|
actions["car_pose"] = [[0.0, 0.0, 0.0] for _ in actions["follow1_pos"]]
|
|
# The checkpoint-1 normalizer records the gripper in the robot's 0..4.5
|
|
# units. Invalid negative values can make Turtle2 reject the whole command
|
|
# packet, so constrain only this scalar channel before sending.
|
|
gripper = np.asarray(actions["follow2_pos"], dtype=np.float64)
|
|
if gripper_min is not None and gripper_max is not None:
|
|
if not np.isfinite([gripper_min, gripper_max]).all() or gripper_min > gripper_max:
|
|
raise ValueError("invalid gripper limits")
|
|
before = gripper[:, 6].copy()
|
|
gripper[:, 6] = np.clip(gripper[:, 6], float(gripper_min), float(gripper_max))
|
|
if not np.array_equal(before, gripper[:, 6]):
|
|
logger.warning(
|
|
"clipped gripper command range from [%.5g, %.5g] to [%.5g, %.5g]",
|
|
float(np.min(before)), float(np.max(before)),
|
|
float(np.min(gripper[:, 6])), float(np.max(gripper[:, 6])),
|
|
)
|
|
actions["follow2_pos"] = gripper.tolist()
|
|
|
|
if state_follow2_pos is not None:
|
|
current = np.asarray(state_follow2_pos, dtype=np.float64).reshape(7)
|
|
traj = np.asarray(actions["follow2_pos"], dtype=np.float64)
|
|
if clip_action_delta:
|
|
# Limit the complete packet relative to the observed pose while
|
|
# preserving the model trajectory's shape. Per-sample cumulative
|
|
# clipping creates a long synthetic ramp toward an unreachable
|
|
# pose and was the source of the repeatable fixed-pose stop.
|
|
clipped, scales = _limit_action_packet(
|
|
traj, current, max_position_delta, max_rotation_delta
|
|
)
|
|
if scales["position"][1] < 1.0 or scales["rotation"][1] < 1.0:
|
|
logger.debug(
|
|
"scaled right-arm packet: position peak %.5g scale %.5g; "
|
|
"rotation peak %.5g scale %.5g",
|
|
scales["position"][0], scales["position"][1],
|
|
scales["rotation"][0], scales["rotation"][1],
|
|
)
|
|
actions["follow2_pos"] = clipped.tolist()
|
|
traj = clipped
|
|
first_delta = np.abs(traj[0, :3] - current[:3])
|
|
limit_eps = 1e-6
|
|
if np.any(first_delta > max_position_delta + limit_eps):
|
|
raise ValueError(
|
|
f"right-arm first position delta {first_delta.tolist()} exceeds "
|
|
f"limit {max_position_delta} m"
|
|
)
|
|
first_rot = np.abs(traj[0, 3:6] - current[3:6])
|
|
if np.any(first_rot > max_rotation_delta + limit_eps):
|
|
raise ValueError(
|
|
f"right-arm first rotation delta {first_rot.tolist()} exceeds "
|
|
f"limit {max_rotation_delta} rad"
|
|
)
|
|
if len(traj) > 1:
|
|
step_pos = np.max(np.abs(np.diff(traj[:, :3], axis=0)), axis=1)
|
|
step_rot = np.max(np.abs(np.diff(traj[:, 3:6], axis=0)), axis=1)
|
|
if np.any(step_pos > max_position_delta + limit_eps) or np.any(step_rot > max_rotation_delta + limit_eps):
|
|
if clip_action_delta:
|
|
# A trajectory can contain Euler-angle wrap discontinuities
|
|
# after preprocessing. In clip mode keep the connection
|
|
# alive and clamp the offending samples to the prior pose.
|
|
logger.warning(
|
|
"right-arm trajectory still has over-limit step after clipping; "
|
|
"holding offending samples"
|
|
)
|
|
safe = traj.copy()
|
|
for i in range(1, len(safe)):
|
|
dp = safe[i, :3] - safe[i - 1, :3]
|
|
dr = safe[i, 3:6] - safe[i - 1, 3:6]
|
|
if np.any(np.abs(dp) > max_position_delta) or np.any(
|
|
np.abs(dr) > max_rotation_delta
|
|
):
|
|
safe[i] = safe[i - 1]
|
|
actions["follow2_pos"] = safe.tolist()
|
|
else:
|
|
raise ValueError("right-arm trajectory contains an over-limit step")
|
|
return actions
|
|
|
|
|
|
async def _handle_client(
|
|
reader,
|
|
writer,
|
|
ws_url,
|
|
instruction,
|
|
allow_send,
|
|
allow_constant_fallbacks,
|
|
action_horizon,
|
|
action_end_ratio,
|
|
action_interpolate_multiplier,
|
|
allow_base_motion,
|
|
hold_left_arm,
|
|
fixed_car_pose,
|
|
fixed_lift,
|
|
max_action_cycles,
|
|
require_right_feedback,
|
|
clip_action_delta,
|
|
max_position_delta,
|
|
max_rotation_delta,
|
|
gripper_min=None,
|
|
gripper_max=None,
|
|
):
|
|
addr = writer.get_extra_info("peername")
|
|
logger.info("Robot connected: %s", addr)
|
|
|
|
async with websockets.connect(ws_url, max_size=None) as ws:
|
|
await ws.recv() # consume server metadata
|
|
|
|
# One observation per connection in dry-run; continuous loop when sending.
|
|
cycles = 0
|
|
previous_right_plan = None
|
|
stop_after_feedback = False
|
|
while True:
|
|
state = await _recv_state(reader)
|
|
if state is None:
|
|
logger.info("Robot closed the connection")
|
|
return
|
|
if previous_right_plan is not None:
|
|
feedback = assess_right_arm_feedback(
|
|
previous_right_plan["initial"],
|
|
state["follow2_pos"],
|
|
previous_right_plan["target"],
|
|
)
|
|
logger.info(
|
|
"right-arm feedback: observed_delta=%s commanded_delta=%s",
|
|
feedback["observed_delta"].tolist(),
|
|
feedback["commanded_delta"].tolist(),
|
|
)
|
|
if feedback["missing_feedback"]:
|
|
logger.warning(
|
|
"right-arm feedback did not change after the prior packet; "
|
|
"bridge sent follow2_pos but the robot did not report EE motion. "
|
|
"Check /follow_pos_cmd_2 subscribers, controller enable state, "
|
|
"/follow2_pos_back, and /joint_information2 on arm-pc."
|
|
)
|
|
if require_right_feedback:
|
|
logger.error(
|
|
"--require-right-feedback set; "
|
|
"stopping before another action packet"
|
|
)
|
|
return
|
|
if stop_after_feedback:
|
|
logger.info("right-arm feedback check complete; stopping after one action packet")
|
|
return
|
|
previous_right_plan = None
|
|
# Turtle2 does not transmit velocity_decomposed. The official
|
|
# Turtle preprocessor accepts the field, so provide a neutral
|
|
# initial value rather than failing on a missing state key. This
|
|
# is deliberately conservative until recurrent velocity semantics
|
|
# are confirmed against a real checkpoint/robot run.
|
|
state.setdefault("velocity_decomposed", [0.0, 0.0, 0.0])
|
|
|
|
images = [await _recv_image(reader, i) for i in range(len(LEGACY_CAM_ORDER))]
|
|
left_img, front_img, right_img = images
|
|
logger.info(
|
|
"recv state keys=%s images=%sx%s",
|
|
list(state.keys()),
|
|
front_img.shape[:2],
|
|
right_img.shape[:2],
|
|
)
|
|
|
|
obs = {
|
|
"state": dict(state),
|
|
"views": {
|
|
SERVE_CAM_KEYS[0]: _img_to_b64(left_img),
|
|
SERVE_CAM_KEYS[1]: _img_to_b64(front_img),
|
|
SERVE_CAM_KEYS[2]: _img_to_b64(right_img),
|
|
},
|
|
"instruction": instruction,
|
|
"infer_mode": "flow",
|
|
}
|
|
logger.info("forwarding to serving: %s", ws_url)
|
|
await ws.send(msgpack.packb(obs))
|
|
raw_resp = await ws.recv()
|
|
if isinstance(raw_resp, str):
|
|
# The official server sends a traceback as a text frame before
|
|
# closing when inference fails. Preserve that diagnostic
|
|
# instead of masking it with msgpack's bytes-only error.
|
|
raise RuntimeError(
|
|
"serving returned a text error frame:\n" + raw_resp
|
|
)
|
|
resp = msgpack.unpackb(raw_resp)
|
|
|
|
actions = prepare_robot_actions(
|
|
resp,
|
|
state_follow1_pos=state.get("follow1_pos"),
|
|
state_follow2_pos=state.get("follow2_pos"),
|
|
state_head_pos=state.get("head_pos"),
|
|
state_lift=state.get("lift"),
|
|
state_car_pose=state.get("car_pose"),
|
|
action_horizon=action_horizon,
|
|
action_end_ratio=action_end_ratio,
|
|
action_interpolate_multiplier=action_interpolate_multiplier,
|
|
allow_base_motion=allow_base_motion,
|
|
hold_left_arm=hold_left_arm,
|
|
fixed_car_pose=fixed_car_pose,
|
|
fixed_lift=fixed_lift,
|
|
allow_constant_fallbacks=allow_constant_fallbacks,
|
|
clip_action_delta=clip_action_delta,
|
|
max_position_delta=max_position_delta,
|
|
max_rotation_delta=max_rotation_delta,
|
|
gripper_min=gripper_min,
|
|
gripper_max=gripper_max,
|
|
)
|
|
logger.info(
|
|
"predicted follow1 T=%d last=%s", len(actions["follow1_pos"]), actions["follow1_pos"][-1]
|
|
)
|
|
logger.info(
|
|
"predicted follow2 T=%d last=%s", len(actions["follow2_pos"]), actions["follow2_pos"][-1]
|
|
)
|
|
try:
|
|
current_right = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7)
|
|
first_right = np.asarray(actions["follow2_pos"][0], dtype=np.float64)
|
|
last_right = np.asarray(actions["follow2_pos"][-1], dtype=np.float64)
|
|
logger.info(
|
|
"right-arm current=%s first=%s first_delta=%s last_delta=%s",
|
|
current_right.tolist(), first_right.tolist(),
|
|
(first_right - current_right).tolist(),
|
|
(last_right - current_right).tolist(),
|
|
)
|
|
except (KeyError, ValueError):
|
|
logger.warning("could not compute right-arm current-to-first delta")
|
|
|
|
if not allow_send:
|
|
logger.warning("DRY-RUN: action NOT sent to robot")
|
|
return
|
|
if not ACTION_SEMANTICS_CONFIRMED:
|
|
logger.error("ACTION_SEMANTICS_CONFIRMED=False; refusing to send to robot")
|
|
return
|
|
|
|
logger.info(
|
|
"wire action T=%d left[first,last]=%s/%s right[first,last]=%s/%s "
|
|
"lift=%s car_pose=%s",
|
|
len(actions["follow1_pos"]),
|
|
actions["follow1_pos"][0], actions["follow1_pos"][-1],
|
|
actions["follow2_pos"][0], actions["follow2_pos"][-1],
|
|
actions["lift"][0], actions["car_pose"][0],
|
|
)
|
|
payload = json.dumps(actions).encode("utf-8")
|
|
writer.write(struct.pack("<L", len(payload)))
|
|
writer.write(payload)
|
|
await writer.drain()
|
|
previous_right_plan = {
|
|
"initial": np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7),
|
|
"target": np.asarray(actions["follow2_pos"][-1], dtype=np.float64).reshape(7),
|
|
}
|
|
logger.info("sent action back to robot T=%d", len(actions["follow1_pos"]))
|
|
cycles += 1
|
|
if max_action_cycles > 0 and cycles >= max_action_cycles:
|
|
if require_right_feedback:
|
|
stop_after_feedback = True
|
|
logger.info(
|
|
"waiting for one robot feedback state before stopping after %d action cycle(s)",
|
|
cycles,
|
|
)
|
|
continue
|
|
logger.warning("stopping after %d action cycle(s)", cycles)
|
|
return
|
|
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--tcp-host", default="0.0.0.0")
|
|
parser.add_argument("--tcp-port", type=int, default=30123)
|
|
parser.add_argument("--ws-url", default="ws://127.0.0.1:32195")
|
|
parser.add_argument("--instruction", default="pick up the water bottel on the chair")
|
|
parser.add_argument(
|
|
"--action-horizon",
|
|
type=int,
|
|
default=32,
|
|
help="Model action horizon (Wall-OSS-0.5 default: 32).",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-send",
|
|
action="store_true",
|
|
help="Request sending actions back to robot (still gated by ACTION_SEMANTICS_CONFIRMED).",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-constant-fallbacks",
|
|
action="store_true",
|
|
help=(
|
|
"Use neutral head/lift/base trajectories when the serving response "
|
|
"omits them. Intended only for protocol dry-runs."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--action-end-ratio",
|
|
type=float,
|
|
default=0.2,
|
|
help="Legacy action trim ratio (default: 0.2).",
|
|
)
|
|
parser.add_argument(
|
|
"--action-interpolate-multiplier",
|
|
type=int,
|
|
default=32,
|
|
help="Legacy linear interpolation multiplier (default: 32).",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-base-motion",
|
|
action="store_true",
|
|
help="Forward model-predicted car_pose instead of holding current pose.",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-left-arm-motion",
|
|
action="store_true",
|
|
help="Forward serializer-provided left-arm commands. Disabled by default for this right-arm-only checkpoint.",
|
|
)
|
|
parser.add_argument(
|
|
"--fixed-car-pose",
|
|
type=float,
|
|
nargs=3,
|
|
metavar=("X", "Y", "YAW"),
|
|
default=None,
|
|
help=(
|
|
"Override every output car_pose with [X, Y, YAW]. "
|
|
"Use only after confirming Turtle2 relative-target semantics."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--fixed-lift",
|
|
type=float,
|
|
default=None,
|
|
help="Force lift target for every action frame (meters, range 0.0-0.47).",
|
|
)
|
|
parser.add_argument(
|
|
"--max-action-cycles", type=int, default=1,
|
|
help="Maximum action packets per connection; 0 means unlimited.",
|
|
)
|
|
parser.add_argument(
|
|
"--require-right-feedback",
|
|
action="store_true",
|
|
help="Stop before the next packet if the previous right-arm command produced no observed EE motion.",
|
|
)
|
|
parser.add_argument(
|
|
"--clip-action-delta", action="store_true",
|
|
help=(
|
|
"Scale each right-arm action packet to safety limits instead of "
|
|
"rejecting it; limits apply to packet excursion and inter-sample steps."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--max-position-delta", type=float, default=0.10,
|
|
help="Maximum right-arm position excursion per inference packet (m).",
|
|
)
|
|
parser.add_argument(
|
|
"--max-rotation-delta", type=float, default=0.50,
|
|
help="Maximum right-arm Euler rotation excursion per inference packet (rad).",
|
|
)
|
|
parser.add_argument(
|
|
"--gripper-min", type=float, default=None,
|
|
help="Optional gripper lower bound; disabled by default.",
|
|
)
|
|
parser.add_argument(
|
|
"--gripper-max", type=float, default=None,
|
|
help="Optional gripper upper bound; disabled by default.",
|
|
)
|
|
parser.add_argument("--log-level", default="INFO")
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=getattr(logging, args.log_level.upper()),
|
|
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
|
)
|
|
|
|
async def client_connected(reader, writer):
|
|
try:
|
|
await _handle_client(
|
|
reader,
|
|
writer,
|
|
args.ws_url,
|
|
args.instruction,
|
|
args.allow_send,
|
|
args.allow_constant_fallbacks,
|
|
args.action_horizon,
|
|
args.action_end_ratio,
|
|
args.action_interpolate_multiplier,
|
|
args.allow_base_motion,
|
|
not args.allow_left_arm_motion,
|
|
args.fixed_car_pose,
|
|
args.fixed_lift,
|
|
args.max_action_cycles,
|
|
args.require_right_feedback,
|
|
args.clip_action_delta,
|
|
args.max_position_delta,
|
|
args.max_rotation_delta,
|
|
args.gripper_min,
|
|
args.gripper_max,
|
|
)
|
|
except Exception as exc:
|
|
logger.error("handler error: %s: %s", type(exc).__name__, exc)
|
|
finally:
|
|
try:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
except Exception:
|
|
pass
|
|
logger.info("robot session closed")
|
|
|
|
server = await asyncio.start_server(client_connected, args.tcp_host, args.tcp_port)
|
|
logger.info(
|
|
"bridge listening on %s:%s -> %s (allow_send=%s)",
|
|
args.tcp_host, args.tcp_port, args.ws_url, args.allow_send,
|
|
)
|
|
async with server:
|
|
await server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(asyncio.run(main()))
|
|
except KeyboardInterrupt:
|
|
sys.exit(0)
|