"""Adapter that maps UDP loopback events onto the existing FSM control flag.""" from __future__ import annotations import queue import socket import threading import time from pathlib import Path from typing import Dict, Optional import yaml from common.joystick import ControlFlag try: from .protocol import InputEnvelope, MotionFrame, format_motion_frame except ImportError: # pragma: no cover - direct script execution fallback from protocol import InputEnvelope, MotionFrame, format_motion_frame class UDPFSMFlag(ControlFlag): """FSM-facing flag produced from decoded UDP events.""" def __init__(self) -> None: super().__init__() self.x_speed_command: float = 0.0 self.y_speed_command: float = 0.0 self.yaw_speed_command: float = 0.0 self.height_cmd: float = 0.89 class UDPFSMController: """Receive UDP packets and expose them as a ControlFlag-like interface.""" def __init__(self) -> None: self.config: Dict[str, object] = {} self.data_mutex = threading.Lock() self._load_config() self._init_data_structures() self.recv_running = False self.recv_thread: Optional[threading.Thread] = None self.rx_socket: Optional[socket.socket] = None def _load_config(self) -> None: config_path = Path(__file__).resolve().parent / "config" / "udp_loopback.yaml" with config_path.open("r", encoding="utf-8") as file: self.config = yaml.safe_load(file) or {} receiver_cfg = self.config.get("receiver", {}) motion_cfg = self.config.get("motion", {}) self.listen_host = receiver_cfg.get("listen_host", "127.0.0.1") self.listen_port = int(receiver_cfg.get("listen_port", 31000)) self.initial_lift = float(motion_cfg.get("initial_lift", 0.89)) self.lift_step = float(motion_cfg.get("lift_step", 0.05)) self.max_surge = float(motion_cfg.get("max_surge", 1.0)) self.max_sway = float(motion_cfg.get("max_sway", 0.5)) self.max_spin = float(motion_cfg.get("max_spin", 0.5)) self.max_lift = float(motion_cfg.get("max_lift", 0.90)) self.min_lift = float(motion_cfg.get("min_lift", 0.65)) self.surge_step = float(motion_cfg.get("surge_step", 0.1)) self.sway_step = float(motion_cfg.get("sway_step", 0.1)) self.spin_step = float(motion_cfg.get("spin_step", 0.1)) def _init_data_structures(self) -> None: self.packet_queue: queue.Queue[InputEnvelope] = queue.Queue(maxsize=128) self.motion_frame = MotionFrame(lift_goal=self.initial_lift) self.udp_flag = UDPFSMFlag() self.udp_flag.height_cmd = self.initial_lift self.last_seq_id = -1 def start(self) -> None: self.rx_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) self.rx_socket.bind((self.listen_host, self.listen_port)) self.rx_socket.settimeout(0.2) self.recv_running = True self.recv_thread = threading.Thread(target=self._recv_loop, daemon=True) self.recv_thread.start() print( f"UDP FSM controller listening on {self.listen_host}:{self.listen_port}" ) def stop(self) -> None: self.recv_running = False if self.recv_thread and self.recv_thread.is_alive(): self.recv_thread.join(timeout=1.0) if self.rx_socket is not None: self.rx_socket.close() self.rx_socket = None print("UDP FSM controller stopped") def _recv_loop(self) -> None: while self.recv_running: try: payload, _ = self.rx_socket.recvfrom(4096) except socket.timeout: continue except OSError: break try: packet = InputEnvelope.decode(payload) except (ValueError, KeyError, TypeError, UnicodeDecodeError) as exc: print(f"[udp_fsm] drop invalid payload: {exc}") continue try: self.packet_queue.put_nowait(packet) except queue.Full: try: self.packet_queue.get_nowait() self.packet_queue.put_nowait(packet) except queue.Empty: pass def update_flag(self) -> None: while not self.packet_queue.empty(): try: packet = self.packet_queue.get_nowait() except queue.Empty: break self._apply_packet(packet) def get_udp_flag(self) -> UDPFSMFlag: with self.data_mutex: flag_copy = UDPFSMFlag() flag_copy.__dict__.update(self.udp_flag.__dict__) return flag_copy def init(self) -> int: print("UDP FSM controller initialized") return 0 def _apply_packet(self, packet: InputEnvelope) -> None: with self.data_mutex: self.last_seq_id = packet.seq_id self.motion_frame.last_event_code = packet.event_code self.motion_frame.last_rx_time = packet.sent_at event_code = packet.event_code if event_code == "pose_home": self.motion_frame.mode_tag = "pose_home" elif event_code == "pose_hold": self.motion_frame.mode_tag = "pose_hold" elif event_code == "mode_stride": self.motion_frame.mode_tag = "mode_stride" elif event_code == "mode_dash": self.motion_frame.mode_tag = "mode_dash" elif event_code == "mode_xrun": self.motion_frame.mode_tag = "mode_xrun" elif event_code == "surge_up": self.motion_frame.surge_goal = min( self.max_surge, self.motion_frame.surge_goal + self.surge_step ) elif event_code == "surge_down": self.motion_frame.surge_goal = max( -self.max_surge, self.motion_frame.surge_goal - self.surge_step ) elif event_code == "sway_left": self.motion_frame.sway_goal = max( -self.max_sway, self.motion_frame.sway_goal - self.sway_step ) elif event_code == "sway_right": self.motion_frame.sway_goal = min( self.max_sway, self.motion_frame.sway_goal + self.sway_step ) elif event_code == "spin_left": self.motion_frame.spin_goal = max( -self.max_spin, self.motion_frame.spin_goal - self.spin_step ) elif event_code == "spin_right": self.motion_frame.spin_goal = min( self.max_spin, self.motion_frame.spin_goal + self.spin_step ) elif event_code == "set_surge": self.motion_frame.surge_goal = max( -self.max_surge, min(self.max_surge, packet.drive_value) ) elif event_code == "set_sway": self.motion_frame.sway_goal = max( -self.max_sway, min(self.max_sway, packet.drive_value) ) elif event_code == "set_spin": self.motion_frame.spin_goal = max( -self.max_spin, min(self.max_spin, packet.drive_value) ) elif event_code == "set_lift": self.motion_frame.lift_goal = max( self.min_lift, min(self.max_lift, packet.drive_value) ) elif event_code == "lift_up": self.motion_frame.lift_goal = min( self.max_lift, self.motion_frame.lift_goal + self.lift_step ) elif event_code == "lift_down": self.motion_frame.lift_goal = max( self.min_lift, self.motion_frame.lift_goal - self.lift_step ) elif event_code == "trim_reset": self.motion_frame.surge_goal = 0.0 self.motion_frame.sway_goal = 0.0 self.motion_frame.spin_goal = 0.0 elif event_code == "session_quit": self.motion_frame.relay_on = False self.motion_frame.mode_tag = "pose_hold" self.recv_running = False self._sync_motion_frame_to_flag() print( f"[udp_fsm] seq={packet.seq_id} key={packet.key_name} " f"{format_motion_frame(self.motion_frame)} " f"fsm={self.udp_flag.fsm_state_command}" ) def _sync_motion_frame_to_flag(self) -> None: mode_to_fsm_command = { "pose_home": "gotoZERO", "pose_hold": "gotoSTOP", "mode_stride": "gotoWALKAMP", "mode_dash": "gotoMYPOLICY", "mode_xrun": "gotoXSIMRUN", } self.udp_flag.enable = self.motion_frame.relay_on self.udp_flag.fsm_state_command = mode_to_fsm_command.get( self.motion_frame.mode_tag, self.udp_flag.fsm_state_command ) self.udp_flag.x_speed_command = self.motion_frame.surge_goal self.udp_flag.y_speed_command = self.motion_frame.sway_goal self.udp_flag.yaw_speed_command = self.motion_frame.spin_goal self.udp_flag.height_cmd = self.motion_frame.lift_goal