242 lines
9.4 KiB
Python
242 lines
9.4 KiB
Python
"""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
|
|
self.last_fsm_command_time = 0.0
|
|
|
|
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 get_last_input_time(self) -> float:
|
|
with self.data_mutex:
|
|
return self.motion_frame.last_rx_time
|
|
|
|
def get_last_fsm_command_time(self) -> float:
|
|
with self.data_mutex:
|
|
return self.last_fsm_command_time
|
|
|
|
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"
|
|
self.last_fsm_command_time = packet.sent_at
|
|
elif event_code == "pose_hold":
|
|
self.motion_frame.mode_tag = "pose_hold"
|
|
self.last_fsm_command_time = packet.sent_at
|
|
elif event_code == "mode_stride":
|
|
self.motion_frame.mode_tag = "mode_stride"
|
|
self.last_fsm_command_time = packet.sent_at
|
|
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",
|
|
}
|
|
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
|