Files
TG3/tg3_omnisocket_transport/omnisocket_xtele_sender.py

596 lines
23 KiB
Python

#!/usr/bin/env python3
"""Session-gate local xTELE data and send it to the robot through OmniSocket.
The raw/local-data stream remains authoritative for arm positions, physical
buttons and hardware diagnostics. When configured, only the processed hand
target from xTELE's command stream is merged into that raw frame. This keeps
xTELE's built-in combo/gesture state machine and BrainCoRevo2 poses without
letting unrelated UI, walking or base commands reach the dual-arm bridge.
The physical left-Z + right-C hold is evaluated locally: idle frames never
leave EAI, while each active interval carries a unique START/ACTIVE/STOP ID.
"""
from __future__ import annotations
import argparse
import copy
import json
import math
import os
from pathlib import Path
import signal
import struct
import time
from typing import Any
import uuid
import zmq
from omnisocket import CONTROL_DEFAULTS, MSG_TYPE_ERROR, Session
MAGIC = b"TG3A"
HEADER = struct.Struct("!4sQQI")
MAX_PAYLOAD_BYTES = 1_048_576
TELEOP_PROTOCOL_VERSION = 2
class XteleSender:
def __init__(self, args: argparse.Namespace) -> None:
if args.start_stop_hold_s <= 0.0:
raise ValueError("start/stop hold time must be positive")
if args.start_marker_frames <= 0:
raise ValueError("start marker frame count must be positive")
if args.source_timeout_s <= 0.0:
raise ValueError("source timeout must be positive")
if args.max_feedback_age_ms <= 0.0:
raise ValueError("maximum KCP feedback age must be positive")
if args.max_pending_frames <= 0:
raise ValueError("maximum pending frame count must be positive")
self.args = args
self.stop = False
self.session: Session | None = None
self.session_connected_at = 0.0
self.started_at = time.time()
self.last_status_at = 0.0
self.last_frame_at = 0.0
self.last_source_frame_at = 0.0
self.last_command_at = 0.0
self.last_error = ""
self.teleop_active = False
self.teleop_session_id: str | None = None
self.teleop_session_seq = 0
self.packet_sequence = time.time_ns()
self.start_markers_remaining = 0
self.combo_started_at: float | None = None
# A service restart must never turn an already-held combo into a start.
self.require_combo_release = True
self.counters = {
"connected": 0,
"reconnects": 0,
"frames_received": 0,
"frames_sent": 0,
"bytes_received": 0,
"bytes_sent": 0,
"dropped_malformed": 0,
"command_frames_received": 0,
"command_frames_accepted": 0,
"command_frames_malformed": 0,
"command_hand_merges": 0,
"remote_errors": 0,
"frames_suppressed_inactive": 0,
"teleop_starts": 0,
"teleop_stops": 0,
"teleop_aborts": 0,
}
signal.signal(signal.SIGINT, self._request_stop)
signal.signal(signal.SIGTERM, self._request_stop)
def _request_stop(self, _signum: int, _frame: Any) -> None:
self.stop = True
def connect(self) -> bool:
self.close_session()
session = Session()
try:
session.connect(
server_addr=self.args.server,
peer_id=self.args.peer_id,
**CONTROL_DEFAULTS,
)
except OSError as exc:
self.last_error = f"OmniSocket connect failed: {exc}"
session.close()
self.write_status(force=True)
return False
self.session = session
self.session_connected_at = time.monotonic()
self.counters["connected"] = 1
self.counters["reconnects"] += 1
self.last_error = ""
self.write_status(force=True)
return True
def close_session(self) -> None:
if self.session is not None:
try:
self.session.close()
except OSError:
pass
self.session = None
self.session_connected_at = 0.0
self.counters["connected"] = 0
def write_status(self, force: bool = False) -> None:
now = time.monotonic()
if not force and now - self.last_status_at < 0.5:
return
session_stats: dict[str, object] = {}
kcp_stats: dict[str, object] = {}
if self.session is not None:
try:
session_stats = self.session.stats()
kcp_stats = self.session.kcp_stats()
except OSError as exc:
self.last_error = f"OmniSocket stats failed: {exc}"
status = {
"role": "xtele_sender",
"server": self.args.server,
"peer_id": self.args.peer_id,
"target_peer": self.args.target_peer,
"zmq_endpoint": self.args.zmq_endpoint,
"cmd_zmq_endpoint": self.args.cmd_zmq_endpoint or None,
"uptime_s": round(time.time() - self.started_at, 1),
"last_frame_age_s": None
if self.last_frame_at == 0.0
else round(time.monotonic() - self.last_frame_at, 4),
"last_command_age_s": None
if self.last_command_at == 0.0
else round(time.monotonic() - self.last_command_at, 4),
"last_source_frame_age_s": None
if self.last_source_frame_at == 0.0
else round(time.monotonic() - self.last_source_frame_at, 4),
"teleop_active": self.teleop_active,
"teleop_session_id": self.teleop_session_id,
"teleop_session_seq": self.teleop_session_seq,
"teleop_combo_hold_s": 0.0
if self.combo_started_at is None
else round(time.monotonic() - self.combo_started_at, 2),
"teleop_require_combo_release": self.require_combo_release,
"application_data_sending": (
self.teleop_active and self.session is not None
),
"last_error": self.last_error,
"session_stats": session_stats,
"kcp_stats": kcp_stats,
**self.counters,
"updated_unix_s": time.time(),
}
path = Path(self.args.status_file)
try:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(json.dumps(status, indent=2) + "\n")
os.replace(temporary, path)
except OSError:
pass
self.last_status_at = now
def _drain_responses(self) -> None:
if self.session is None:
return
try:
while True:
response = self.session.recv(timeout_ms=0)
if response is None:
break
_from_peer, msg_type, payload = response
if msg_type == MSG_TYPE_ERROR:
self.counters["remote_errors"] += 1
detail = payload.decode("utf-8", errors="replace")
self._abort_teleop(f"OmniSocket remote error: {detail}")
self.close_session()
return
except OSError as exc:
self._abort_teleop(f"OmniSocket receive failed: {exc}")
self.close_session()
@staticmethod
def _valid_hand_side(value: object) -> bool:
if isinstance(value, bool):
return False
if isinstance(value, (int, float)):
numeric = float(value)
return math.isfinite(numeric) and -0.05 <= numeric <= 1.05
if not isinstance(value, list) or len(value) != 6:
return False
try:
numeric_values = [float(item) for item in value]
except (TypeError, ValueError):
return False
return all(
math.isfinite(item) and -0.05 <= item <= 1.05
for item in numeric_values
)
@classmethod
def _processed_hand_position(
cls, command: dict[str, object]
) -> dict[str, object] | None:
try:
position = command["hand"]["position"] # type: ignore[index]
left = position["left"] # type: ignore[index]
right = position["right"] # type: ignore[index]
except (KeyError, TypeError):
return None
if not cls._valid_hand_side(left) or not cls._valid_hand_side(right):
return None
return {"left": copy.deepcopy(left), "right": copy.deepcopy(right)}
@classmethod
def _build_payload(
cls,
data: dict[str, object],
command: dict[str, object] | None,
session_id: str,
session_seq: int,
session_state: str,
stop_reason: str = "",
) -> tuple[bytes, bool]:
merged = False
try:
position = cls._processed_hand_position(command) if command else None
hand = data.get("hand")
if position is not None and isinstance(hand, dict):
hand["position"] = position
merged = True
metadata = data.get("tg3_transport")
if not isinstance(metadata, dict):
metadata = {}
data["tg3_transport"] = metadata
if merged:
metadata["processed_hand_from_xtele_cmd"] = True
assert command is not None
metadata["xtele_cmd_timestamp"] = command.get("timestamp")
else:
metadata.pop("processed_hand_from_xtele_cmd", None)
metadata.pop("xtele_cmd_timestamp", None)
# Always overwrite untrusted source metadata. The robot accepts
# start/active/stop only from this sender and expected Omni peer.
metadata["protocol_version"] = TELEOP_PROTOCOL_VERSION
metadata["session_id"] = session_id
metadata["session_seq"] = session_seq
metadata["session_state"] = session_state
if stop_reason:
metadata["stop_reason"] = stop_reason
else:
metadata.pop("stop_reason", None)
encoded = json.dumps(
data, ensure_ascii=False, separators=(",", ":")
).encode("utf-8")
except (TypeError, ValueError):
raise ValueError("cannot encode xTELE session payload")
return encoded, merged
@staticmethod
def _start_stop_pressed(data: dict[str, object]) -> bool:
try:
buttons = data["button"]
left = buttons["left"] # type: ignore[index]
right = buttons["right"] # type: ignore[index]
return (
len(left) >= 3 # type: ignore[arg-type]
and len(right) >= 3 # type: ignore[arg-type]
and bool(left[2]) # type: ignore[index]
and bool(right[2]) # type: ignore[index]
)
except (KeyError, TypeError):
return False
def _update_teleop_gate(
self, now: float, data: dict[str, object]
) -> str | None:
pressed = self._start_stop_pressed(data)
if self.require_combo_release:
self.combo_started_at = None
if not pressed:
self.require_combo_release = False
return None
if not pressed:
self.combo_started_at = None
return None
if self.combo_started_at is None:
self.combo_started_at = now
if now - self.combo_started_at < self.args.start_stop_hold_s:
return None
self.combo_started_at = None
self.require_combo_release = True
if self.teleop_active:
self.teleop_active = False
self.counters["teleop_stops"] += 1
return "stop"
self.teleop_active = True
self.teleop_session_id = uuid.uuid4().hex
self.teleop_session_seq = 0
self.start_markers_remaining = self.args.start_marker_frames
self.counters["teleop_starts"] += 1
return "start"
def _abort_teleop(self, reason: str) -> None:
if self.teleop_active:
self.counters["teleop_aborts"] += 1
self.teleop_active = False
self.teleop_session_id = None
self.teleop_session_seq = 0
self.start_markers_remaining = 0
self.combo_started_at = None
self.require_combo_release = True
self.last_error = reason
def _send_payload(self, payload: bytes) -> bool:
if self.session is None:
self._abort_teleop("OmniSocket session is unavailable")
return False
unhealthy = self._session_unhealthy_reason()
if unhealthy:
self._abort_teleop(unhealthy)
self.close_session()
return False
self.packet_sequence += 1
packet = HEADER.pack(
MAGIC, self.packet_sequence, time.time_ns(), len(payload)
) + payload
try:
self.session.send(to=self.args.target_peer, data=packet)
except OSError as exc:
self._abort_teleop(f"OmniSocket send failed: {exc}")
self.close_session()
return False
self.counters["frames_sent"] += 1
self.counters["bytes_sent"] += len(payload)
self.last_frame_at = time.monotonic()
self.last_error = ""
return True
def _session_unhealthy_reason(self) -> str | None:
if self.session is None:
return "OmniSocket session is unavailable"
try:
session_stats = self.session.stats()
kcp_stats = self.session.kcp_stats()
if int(session_stats.get("connected", 0)) != 1 or int(
session_stats.get("registered", 0)
) != 1:
return "OmniSocket session lost registration"
pending = int(kcp_stats.get("snd_queue", 0)) + int(
kcp_stats.get("snd_buffer", 0)
)
if pending > self.args.max_pending_frames:
return (
f"OmniSocket pending queue reached {pending} frames; "
"session invalidated to prevent stale replay"
)
feedback_age_ms = float(kcp_stats.get("last_feedback_age_ms", 0.0))
connection_age_s = time.monotonic() - self.session_connected_at
if (
connection_age_s >= 0.5
and feedback_age_ms > self.args.max_feedback_age_ms
):
return (
f"OmniSocket feedback stale for {feedback_age_ms:.0f} ms; "
"session invalidated to prevent stale replay"
)
except (OSError, TypeError, ValueError) as exc:
return f"cannot verify OmniSocket session health: {exc}"
return None
def _flush_session(self, timeout_s: float = 0.75) -> None:
"""Bound the final STOP flush, without sending another business frame."""
deadline = time.monotonic() + timeout_s
while self.session is not None and time.monotonic() < deadline:
try:
stats = self.session.kcp_stats()
pending = int(stats.get("snd_queue", 0)) + int(
stats.get("snd_buffer", 0)
)
except (OSError, TypeError, ValueError):
return
if pending == 0:
return
self._drain_responses()
time.sleep(0.01)
def run(self) -> int:
context = zmq.Context()
source = context.socket(zmq.SUB)
source.setsockopt(zmq.SUBSCRIBE, b"")
source.setsockopt(zmq.CONFLATE, 1)
source.setsockopt(zmq.RCVHWM, 1)
source.setsockopt(zmq.LINGER, 0)
source.connect(self.args.zmq_endpoint)
poller = zmq.Poller()
poller.register(source, zmq.POLLIN)
command_source = None
if self.args.cmd_zmq_endpoint:
command_source = context.socket(zmq.SUB)
command_source.setsockopt(zmq.SUBSCRIBE, b"")
command_source.setsockopt(zmq.CONFLATE, 1)
command_source.setsockopt(zmq.RCVHWM, 1)
command_source.setsockopt(zmq.LINGER, 0)
command_source.connect(self.args.cmd_zmq_endpoint)
poller.register(command_source, zmq.POLLIN)
latest_command: dict[str, object] | None = None
try:
while not self.stop:
events = dict(poller.poll(100))
if command_source is not None and command_source in events:
command_raw = command_source.recv()
self.counters["command_frames_received"] += 1
try:
parsed_command = json.loads(command_raw)
if not isinstance(parsed_command, dict):
raise ValueError("xTELE command must be a JSON object")
if self._processed_hand_position(parsed_command) is None:
raise ValueError(
"xTELE command has no valid bilateral hand target"
)
latest_command = parsed_command
self.last_command_at = time.monotonic()
self.counters["command_frames_accepted"] += 1
except (TypeError, ValueError, json.JSONDecodeError):
self.counters["command_frames_malformed"] += 1
if source not in events:
now = time.monotonic()
if (
self.last_source_frame_at != 0.0
and now - self.last_source_frame_at
> self.args.source_timeout_s
):
self.combo_started_at = None
if self.teleop_active:
self._abort_teleop(
"local xTELE source became stale; a new Z+C hold "
"is required"
)
self.close_session()
self._drain_responses()
self.write_status()
continue
raw = source.recv()
self.counters["frames_received"] += 1
self.counters["bytes_received"] += len(raw)
if not raw or len(raw) > MAX_PAYLOAD_BYTES:
self.counters["dropped_malformed"] += 1
continue
try:
data = json.loads(raw)
if not isinstance(data, dict):
raise ValueError("xTELE raw frame must be a JSON object")
except (TypeError, ValueError, json.JSONDecodeError):
# A malformed/frozen frame may never contribute time to a
# physical three-second start/stop hold.
self.combo_started_at = None
self.counters["dropped_malformed"] += 1
continue
now = time.monotonic()
self.last_source_frame_at = now
transition = self._update_teleop_gate(now, data)
command = None
if (
latest_command is not None
and now - self.last_command_at <= self.args.cmd_max_age_s
):
command = latest_command
if not self.teleop_active and transition != "stop":
self.counters["frames_suppressed_inactive"] += 1
self._drain_responses()
self.write_status()
continue
session_id = self.teleop_session_id
if not session_id:
self._abort_teleop("teleoperation session ID is unavailable")
self.counters["frames_suppressed_inactive"] += 1
continue
if self.session is None and not self.connect():
self._abort_teleop(
"cannot start teleoperation because OmniSocket Hub is "
"unavailable; release Z+C before retrying"
)
self.counters["frames_suppressed_inactive"] += 1
continue
self.teleop_session_seq += 1
if transition == "stop":
session_state = "stop"
stop_reason = "operator"
elif self.start_markers_remaining > 0:
session_state = "start"
stop_reason = ""
else:
session_state = "active"
stop_reason = ""
try:
payload, merged = self._build_payload(
data,
command,
session_id,
self.teleop_session_seq,
session_state,
stop_reason,
)
except ValueError:
self.counters["dropped_malformed"] += 1
continue
if merged:
self.counters["command_hand_merges"] += 1
if len(payload) > MAX_PAYLOAD_BYTES:
self.counters["dropped_malformed"] += 1
continue
sent = self._send_payload(payload)
if sent and session_state == "start":
self.start_markers_remaining -= 1
if transition == "stop":
# STOP is the final xTELE business frame. The underlying
# registered OmniSocket session remains warm for low-latency
# next start and KCP delivery, but no arm data follows.
self.teleop_session_id = None
self.teleop_session_seq = 0
self.start_markers_remaining = 0
self._flush_session()
self.close_session()
self._drain_responses()
self.write_status()
finally:
source.close()
if command_source is not None:
command_source.close()
context.term()
self.close_session()
self.write_status(force=True)
return 0
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--server", required=True)
parser.add_argument("--peer-id", required=True)
parser.add_argument("--target-peer", required=True)
parser.add_argument("--zmq-endpoint", required=True)
parser.add_argument(
"--cmd-zmq-endpoint",
default="",
help=(
"optional xTELE processed-command PUB endpoint; only its bilateral "
"hand.position target is merged into the raw frame"
),
)
parser.add_argument("--cmd-max-age-s", type=float, default=0.25)
parser.add_argument("--source-timeout-s", type=float, default=0.25)
parser.add_argument("--max-feedback-age-ms", type=float, default=500.0)
parser.add_argument("--max-pending-frames", type=int, default=100)
parser.add_argument("--start-stop-hold-s", type=float, default=3.0)
parser.add_argument(
"--start-marker-frames",
type=int,
default=50,
help="repeat START for this many source frames before ACTIVE",
)
parser.add_argument("--status-file", required=True)
return parser.parse_args()
if __name__ == "__main__":
raise SystemExit(XteleSender(parse_args()).run())