First Commit

Deploy_Tienkung/udp_loopback/udp_loopback_node.py

@@ -0,0 +1,227 @@
+"""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 == "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 == "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()