"""Receiver-side localhost UDP loopback node with a control loop layout.""" from __future__ import annotations import queue import socket import threading import time from pathlib import Path from typing import Dict, Optional import yaml 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 UDPLoopbackNode: """Standalone UDP receiver that mirrors rl_control_node.py structure.""" def __init__(self, debug: bool = False) -> None: self.debug = debug self.config: Dict[str, object] = {} self._load_config() self._init_data_structures() self._init_udp_interfaces() self._init_control_system() self._start_threads() def _load_config(self) -> None: self.config_file = Path(__file__).resolve().parent / "config" / "udp_loopback.yaml" with self.config_file.open("r", encoding="utf-8") as file: self.config = yaml.safe_load(file) or {} self.dt = float(self.config.get("dt", 0.01)) self.receiver_cfg = self.config.get("receiver", {}) self.motion_cfg = self.config.get("motion", {}) self.listen_host = self.receiver_cfg.get("listen_host", "127.0.0.1") self.listen_port = int(self.receiver_cfg.get("listen_port", 31000)) self.initial_lift = float(self.motion_cfg.get("initial_lift", 0.89)) self.lift_step = float(self.motion_cfg.get("lift_step", 0.05)) self.max_surge = float(self.motion_cfg.get("max_surge", 1.0)) self.max_sway = float(self.motion_cfg.get("max_sway", 0.5)) self.max_spin = float(self.motion_cfg.get("max_spin", 0.5)) self.max_lift = float(self.motion_cfg.get("max_lift", 0.90)) self.min_lift = float(self.motion_cfg.get("min_lift", 0.65)) self.surge_step = float(self.motion_cfg.get("surge_step", 0.1)) self.sway_step = float(self.motion_cfg.get("sway_step", 0.1)) self.spin_step = float(self.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.frame_lock = threading.Lock() self.last_seq_id = -1 def _init_udp_interfaces(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) def _init_control_system(self) -> None: self.control_running = False self.recv_running = False self.control_thread: Optional[threading.Thread] = None self.recv_thread: Optional[threading.Thread] = None def _start_threads(self) -> None: self.recv_running = True self.control_running = True self.recv_thread = threading.Thread(target=self._udp_recv_loop, daemon=True) self.recv_thread.start() self.control_thread = threading.Thread(target=self._control_loop, daemon=True) self.control_thread.start() print( f"UDP loopback node listening on {self.listen_host}:{self.listen_port} " f"with dt={self.dt:.3f}s" ) def _udp_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, yaml.YAMLError) as exc: print(f"[udp_loopback] 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 _control_loop(self) -> None: print("UDP loopback control loop started") while self.control_running: loop_start = time.perf_counter() self._process_udp_packets() self._precise_sleep_until(loop_start + self.dt) print("UDP loopback control loop ended") def _process_udp_packets(self) -> None: while not self.packet_queue.empty(): try: packet = self.packet_queue.get_nowait() except queue.Empty: break self._apply_packet(packet) def _apply_packet(self, packet: InputEnvelope) -> None: with self.frame_lock: 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 == "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 == "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 = "session_quit" self.control_running = False self.recv_running = False print( f"[udp_loopback] seq={packet.seq_id} key={packet.key_name} " f"{format_motion_frame(self.motion_frame)}" ) def _precise_sleep_until(self, target_time: float) -> None: current_time = time.perf_counter() sleep_time = target_time - current_time if sleep_time <= 0: return if sleep_time > 0.003: time.sleep(sleep_time * 0.9) while time.perf_counter() < target_time: pass else: while time.perf_counter() < target_time: pass def destroy(self) -> None: self.control_running = False self.recv_running = False if self.control_thread and self.control_thread.is_alive(): self.control_thread.join(timeout=1.0) if self.recv_thread and self.recv_thread.is_alive(): self.recv_thread.join(timeout=1.0) self.rx_socket.close() def main() -> None: node = UDPLoopbackNode(debug=False) try: while node.control_running: time.sleep(0.2) except KeyboardInterrupt: pass finally: node.destroy() if __name__ == "__main__": main()