"""Robot arm controller — real DDS motor commands and JSONL motion replay. Production-grade replay engine ported from G1_Lootah/Controller/g1_replay_trigger_r2x.py. Features: body-lock, ramp-in/out interpolation, watchdog, speed control, CRC. Falls back gracefully to simulation when the Unitree SDK is unavailable. """ from __future__ import annotations import json import threading import time from dataclasses import dataclass from pathlib import Path from typing import Any from Project.Sanad.config import ( G1_NUM_MOTOR, KD_HIGH, KD_LOW, KD_WRIST, KP_HIGH, KP_LOW, KP_WRIST, MOTIONS_DIR, REPLAY_HZ, WEAK_MOTORS, WRIST_MOTORS, ) from Project.Sanad.core.config_loader import section as _cfg_section from Project.Sanad.core.event_bus import bus from Project.Sanad.core.logger import get_logger log = get_logger("arm_controller") _AC = _cfg_section("motion", "arm_controller") # G1 hardware constants — single source in config/core_config.json from Project.Sanad.config import ENABLE_ARM_SDK_INDEX RAMP_IN_STEPS = _AC.get("ramp_in_steps", 60) # ~1.0s smooth move to start pose RAMP_OUT_STEPS = _AC.get("ramp_out_steps", 180) # ~3.0s smooth return to home WATCHDOG_TIMEOUT = _AC.get("watchdog_timeout_sec", 0.25) # hold last pose if state stale WATCHDOG_DISABLE_AFTER = _AC.get("watchdog_disable_after_sec", 1.0) # abort if state stale this long ARM_INDICES = range( _AC.get("arm_indices_start", 15), _AC.get("arm_indices_stop", 29), ) # -- SDK import (optional) -- try: from unitree_sdk2py.core.channel import ( ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber, ) from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_, LowState_ from unitree_sdk2py.utils.crc import CRC # IDL factory — LowCmd_() with no args fails because the dataclass # has 5 required positional fields. The SDK ships a default factory # named `unitree_hg_msg_dds__LowCmd_` that constructs a fully-zeroed # instance with the right number of motor_cmd entries. try: from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ _make_low_cmd = unitree_hg_msg_dds__LowCmd_ except ImportError: _make_low_cmd = None _HAS_SDK = True except ImportError: _HAS_SDK = False _make_low_cmd = None log.warning("Unitree SDK not available — running in simulation mode") # Optional G1 high-level SDK action client (built-in actions) try: from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient _HAS_LOCO = True except ImportError: _HAS_LOCO = False @dataclass class Action: name: str id: int file: str = "" # JSONL filename (empty = SDK built-in) category: str = "sdk" # "sdk" | "jsonl" # -- SDK actions (fixed — built into Unitree firmware) -- SDK_ACTIONS: list[Action] = [ Action("release_arm", 0, category="sdk"), Action("shake_hand", 1, category="sdk"), Action("high_five", 2, category="sdk"), Action("hug", 3, category="sdk"), Action("high_wave", 4, category="sdk"), Action("clap", 5, category="sdk"), Action("face_wave", 6, category="sdk"), Action("left_kiss", 7, category="sdk"), Action("heart", 8, category="sdk"), Action("right_heart", 9, category="sdk"), Action("hands_up", 10, category="sdk"), Action("x_ray", 11, category="sdk"), Action("right_hand_up", 12, category="sdk"), Action("reject", 13, category="sdk"), Action("right_kiss", 14, category="sdk"), Action("two_hand_kiss", 15, category="sdk"), ] # Next auto-ID for JSONL actions starts after SDK range. _JSONL_ID_START = _AC.get("jsonl_id_start", 100) def _scan_jsonl_actions() -> list[Action]: """Auto-discover all .jsonl files in data/motions/ and create actions. Called at startup and whenever the dashboard requests a refresh. The action name is derived from the filename (without extension), with underscores replacing hyphens/spaces. """ MOTIONS_DIR.mkdir(parents=True, exist_ok=True) actions = [] for idx, path in enumerate(sorted(MOTIONS_DIR.glob("*.jsonl"))): name = path.stem.replace("-", "_").replace(" ", "_") actions.append(Action( name=name, id=_JSONL_ID_START + idx, file=path.name, category="jsonl", )) return actions def rebuild_action_registry() -> tuple[list[Action], dict[int, Action], dict[str, Action]]: """Rebuild the full action list from SDK + disk scan. Called on startup and refresh.""" jsonl_actions = _scan_jsonl_actions() all_actions = list(SDK_ACTIONS) + jsonl_actions by_id = {a.id: a for a in all_actions} by_name = {a.name: a for a in all_actions} log.info("Action registry: %d SDK + %d JSONL = %d total", len(SDK_ACTIONS), len(jsonl_actions), len(all_actions)) return all_actions, by_id, by_name # Initial build ACTIONS, ACTION_BY_ID, ACTION_BY_NAME = rebuild_action_registry() def _lerp_q(q_start: list[float], q_end: list[float], t: float) -> list[float]: """Linear interpolation between two joint-position vectors, t in [0,1].""" return [a + (b - a) * t for a, b in zip(q_start, q_end)] def _load_frames(path: Path) -> list[dict[str, Any]]: """Read JSONL file, return list of frames with 't' and 'q' keys.""" frames = [] with open(path, "r") as f: for line in f: line = line.strip() if not line: continue data = json.loads(line) if "q" in data: frames.append(data) return frames def _load_home_q(home_file: str = "arm_home.jsonl") -> list[float] | None: path = MOTIONS_DIR / home_file if not path.exists(): return None frames = _load_frames(path) return frames[0]["q"] if frames else None class ArmController: """Thread-safe arm controller with real DDS replay and simulation fallback.""" def __init__(self): self._lock = threading.Lock() self._cancel = threading.Event() self._is_busy = False self._last_action_time = 0.0 self.cooldown_sec = 1.0 self._initialized = False # DDS handles (set in init()) self._arm_pub = None self._state_sub = None self._low_cmd = None self._crc = None self._loco_client = None # Live state from LowState_ subscriber self._current_q: list[float] = [0.0] * G1_NUM_MOTOR self._last_state_time = 0.0 self._state_lock = threading.Lock() self._first_state_event = threading.Event() # Cached motion file metadata self._motion_files_cache: dict[str, dict[str, Any]] = {} # -- initialization -- def init(self, network_interface: str = "lo") -> bool: if self._initialized: return True if not _HAS_SDK: log.info("Simulation mode — DDS init skipped") return False try: ChannelFactoryInitialize(0, network_interface) self._arm_pub = ChannelPublisher("rt/arm_sdk", LowCmd_) self._arm_pub.Init() self._state_sub = ChannelSubscriber("rt/lowstate", LowState_) self._state_sub.Init(self._on_low_state, 10) # IDL types need the SDK's default factory — bare LowCmd_() fails # because the dataclass has 5 required positional fields. if _make_low_cmd is not None: self._low_cmd = _make_low_cmd() else: # Last-resort: try a few constructor signatures try: self._low_cmd = LowCmd_() except TypeError: # Build with explicit zeroed fields from unitree_sdk2py.idl.unitree_hg.msg.dds_ import MotorCmd_ try: from unitree_sdk2py.idl.default import ( unitree_hg_msg_dds__MotorCmd_ as _make_motor_cmd, ) except ImportError: _make_motor_cmd = lambda: MotorCmd_( mode=0, q=0.0, dq=0.0, tau=0.0, kp=0.0, kd=0.0, reserve=0, ) self._low_cmd = LowCmd_( mode_pr=0, mode_machine=0, motor_cmd=[_make_motor_cmd() for _ in range(35)], reserve=[0, 0, 0, 0], crc=0, ) self._crc = CRC() # High-level SDK action client (for built-in actions) if _HAS_LOCO: try: self._loco_client = LocoClient() self._loco_client.SetTimeout(10.0) self._loco_client.Init() log.info("LocoClient initialized — built-in SDK actions available") except Exception as exc: log.warning("LocoClient init failed: %s — built-in actions disabled", exc) self._loco_client = None self._initialized = True log.info("DDS initialized on %s", network_interface) except Exception as exc: log.error("DDS init failed: %s", exc) return self._initialized def _on_low_state(self, msg): """Callback from DDS subscriber — updates current joint positions.""" with self._state_lock: self._current_q = [float(msg.motor_state[i].q) for i in range(G1_NUM_MOTOR)] self._last_state_time = time.monotonic() if not self._first_state_event.is_set(): self._first_state_event.set() def wait_for_state(self, timeout: float = 2.0) -> bool: """Block until first LowState_ callback fires (or timeout). Returns True if state received.""" return self._first_state_event.wait(timeout=timeout) # -- internal API exposed for teaching/macro_player (encapsulation boundary) -- def get_current_q(self) -> list[float]: """Public read of current joint positions.""" return self._get_current_q() def send_frame(self, arm_target_q: list[float], body_lock_q: list[float]): """Public single-frame send. Use only inside a controlled playback loop.""" self._send_frame(arm_target_q, body_lock_q) def disable(self): """Public disable — releases arm SDK control.""" self._disable_sdk() def state_age(self) -> float: """Seconds since last LowState_ callback.""" return self._state_age() def _get_current_q(self) -> list[float]: with self._state_lock: return list(self._current_q) def _state_age(self) -> float: with self._state_lock: return time.monotonic() - self._last_state_time if self._last_state_time else 999.0 # -- frame sending (real DDS with CRC) -- def _send_frame(self, arm_target_q: list[float], body_lock_q: list[float]): """Send one motor frame via DDS. Body stays locked, arms get target.""" if not self._initialized or self._low_cmd is None: return # Enable ARM_SDK self._low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 1.0 for i in range(G1_NUM_MOTOR): self._low_cmd.motor_cmd[i].mode = 1 self._low_cmd.motor_cmd[i].dq = 0.0 self._low_cmd.motor_cmd[i].tau = 0.0 # Arms get replay data, body stays locked if i >= 15: self._low_cmd.motor_cmd[i].q = arm_target_q[i] else: self._low_cmd.motor_cmd[i].q = body_lock_q[i] # Per-motor gains if i in WEAK_MOTORS: kp, kd = KP_LOW, KD_LOW elif i in WRIST_MOTORS: kp, kd = KP_WRIST, KD_WRIST else: kp, kd = KP_HIGH, KD_HIGH self._low_cmd.motor_cmd[i].kp = kp self._low_cmd.motor_cmd[i].kd = kd self._low_cmd.crc = self._crc.Crc(self._low_cmd) self._arm_pub.Write(self._low_cmd) def _disable_sdk(self): """Send 5 disable frames to cleanly release arm control.""" if not self._initialized or self._low_cmd is None: return for _ in range(5): self._low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0 self._low_cmd.crc = self._crc.Crc(self._low_cmd) self._arm_pub.Write(self._low_cmd) time.sleep(0.02) # -- replay engine -- def replay_file(self, filepath: str, speed: float = 1.0): """Play a JSONL motion file with full production replay logic. Args: filepath: Path to .jsonl motion file. speed: Playback speed multiplier (1.0 = normal). """ with self._lock: if self._is_busy: log.warning("replay_file: arm busy, skipping") return self._is_busy = True self._cancel.clear() try: self._replay_file_inner(filepath, speed) finally: with self._lock: self._is_busy = False self._last_action_time = time.monotonic() def _replay_file_inner(self, filepath: str, speed: float = 1.0): path = Path(filepath) if not path.is_absolute(): path = MOTIONS_DIR / path if not path.exists(): raise FileNotFoundError(f"Motion file not found: {path}") frames = _load_frames(path) if not frames: log.warning("Empty motion file: %s", path) return home_q = _load_home_q() or [0.0] * G1_NUM_MOTOR if not _HAS_SDK: duration = len(frames) / REPLAY_HZ / speed if speed else len(frames) / REPLAY_HZ log.info("[SIM] Replay %s (%.1fs, %d frames, speed=%.1f)", path.name, duration, len(frames), speed) self._sim_replay(frames, speed) return log.info("Replay %s (%d frames @ %.0f Hz, speed=%.1f)", path.name, len(frames), REPLAY_HZ, speed) # CRITICAL: wait for first LowState_ callback before reading _current_q # Otherwise we'd lock the body to all-zeros and crash the robot. if not self._first_state_event.is_set(): log.warning("Waiting for first LowState message...") if not self._first_state_event.wait(timeout=2.0): log.error("No LowState received in 2s — refusing to replay (would lock body to zeros)") return body_lock_q = self._get_current_q() current_arm_q = list(body_lock_q) start_arm_q = frames[0]["q"] interval = 1.0 / REPLAY_HZ # Phase 1: Ramp-in — interpolate from current pose to first frame log.debug("Ramp-in (%d steps)", RAMP_IN_STEPS) for step in range(RAMP_IN_STEPS): if self._cancel.is_set(): self._return_home(current_arm_q, body_lock_q, home_q) return t = (step + 1) / RAMP_IN_STEPS interp = _lerp_q(current_arm_q, start_arm_q, t) self._send_frame(interp, body_lock_q) time.sleep(interval) # Phase 2: Play frames with timing and watchdog log.debug("Playing %d frames", len(frames)) play_elapsed = 0.0 last_real = time.monotonic() t0_frame = frames[0].get("t", 0.0) for frame in frames: if self._cancel.is_set(): break # Watchdog age = self._state_age() if age > WATCHDOG_DISABLE_AFTER: log.error("Watchdog abort — LowState stale %.2fs", age) self._disable_sdk() return if age > WATCHDOG_TIMEOUT: log.warning("Watchdog hold — LowState stale %.2fs", age) now_real = time.monotonic() play_elapsed += (now_real - last_real) * speed last_real = now_real frame_t = frame.get("t", 0.0) - t0_frame if frame_t > play_elapsed: time.sleep(min(frame_t - play_elapsed, interval)) self._send_frame(frame["q"], body_lock_q) # NOTE: timing is driven by frame_t above — no extra sleep here # Phase 3: Return home last_arm_q = frames[-1]["q"] if frames else current_arm_q self._return_home(last_arm_q, body_lock_q, home_q) def _return_home(self, from_q: list[float], body_lock_q: list[float], home_q: list[float]): """Smooth interpolation back to home pose, then disable SDK.""" log.debug("Returning home (%d steps)", RAMP_OUT_STEPS) interval = 1.0 / REPLAY_HZ for step in range(RAMP_OUT_STEPS): if self._cancel.is_set(): break t = (step + 1) / RAMP_OUT_STEPS interp = _lerp_q(from_q, home_q, t) self._send_frame(interp, body_lock_q) time.sleep(interval) self._disable_sdk() log.info("Home reached, SDK disabled") def _sim_replay(self, frames: list[dict], speed: float): """Simulation replay — emit events, sleep for equivalent duration.""" interval = 1.0 / REPLAY_HZ for i, frame in enumerate(frames): if self._cancel.is_set(): log.info("[SIM] Replay cancelled at frame %d/%d", i, len(frames)) return time.sleep(interval / max(speed, 0.1)) log.info("[SIM] Replay complete") # -- public API -- @property def is_busy(self) -> bool: return self._is_busy def cancel(self): """Cancel the currently running replay.""" self._cancel.set() def refresh_actions(self): """Re-scan data/motions/ and rebuild the action registry.""" global ACTIONS, ACTION_BY_ID, ACTION_BY_NAME ACTIONS, ACTION_BY_ID, ACTION_BY_NAME = rebuild_action_registry() def list_actions(self) -> list[dict[str, Any]]: return [ {"id": a.id, "name": a.name, "file": a.file, "category": a.category} for a in ACTIONS ] def list_motion_files(self) -> list[dict[str, Any]]: """List all JSONL files in data/motions/ with metadata. Caches frame count by (path, mtime) to avoid re-parsing megabytes of JSONL on every dashboard refresh. """ MOTIONS_DIR.mkdir(parents=True, exist_ok=True) result = [] for p in sorted(MOTIONS_DIR.glob("*.jsonl")): stat = p.stat() cache_key = f"{p}:{stat.st_mtime_ns}" cached = self._motion_files_cache.get(cache_key) if cached is None: frames = _load_frames(p) duration = len(frames) / REPLAY_HZ if frames else 0 cached = { "name": p.name, "path": str(p), "frames": len(frames), "duration_sec": round(duration, 2), "size_kb": round(stat.st_size / 1024, 1), } # Drop stale entries for this path before adding new one stale = [k for k in self._motion_files_cache if k.startswith(f"{p}:")] for k in stale: self._motion_files_cache.pop(k, None) self._motion_files_cache[cache_key] = cached result.append(cached) return result def trigger_by_id(self, action_id: int, speed: float = 1.0): action = ACTION_BY_ID.get(action_id) if action is None: raise KeyError(f"Unknown action id: {action_id}") self._execute(action, speed) def trigger_by_name(self, name: str, speed: float = 1.0): action = ACTION_BY_NAME.get(name) if action is None: raise KeyError(f"Unknown action: {name}") self._execute(action, speed) def _execute(self, action: Action, speed: float = 1.0): with self._lock: if self._is_busy: log.warning("Arm busy, skipping %s", action.name) return self._is_busy = True self._cancel.clear() try: bus.emit_sync("motion.action_started", action=action.name) if action.file: self._replay_file_inner(action.file, speed=speed) else: self._run_sdk_action(action) except Exception as exc: log.error("Action %s failed: %s", action.name, exc) finally: with self._lock: self._is_busy = False self._last_action_time = time.monotonic() bus.emit_sync("motion.action_done", action=action.name) def _run_sdk_action(self, action: Action): if not _HAS_SDK: log.info("[SIM] SDK action: %s (id=%d)", action.name, action.id) time.sleep(2.0) return if self._loco_client is None: log.warning( "SDK action %s requested but LocoClient not available — skipping", action.name, ) return log.info("SDK action: %s (id=%d)", action.name, action.id) try: # G1 LocoClient exposes ExecuteAction(id) for built-in actions self._loco_client.ExecuteAction(action.id) # Built-in actions are blocking on the robot side, but the SDK call returns # immediately. Give it a sensible default duration so we don't hammer it. time.sleep(2.5) except Exception as exc: log.error("SDK action %s failed: %s", action.name, exc) def status(self) -> dict[str, Any]: return { "initialized": self._initialized, "sdk_available": _HAS_SDK, "busy": self._is_busy, "state_age_sec": round(self._state_age(), 3), "sdk_actions": len(SDK_ACTIONS), "jsonl_actions": len([a for a in ACTIONS if a.category == "jsonl"]), "total_actions": len(ACTIONS), "total_motion_files": len(list(MOTIONS_DIR.glob("*.jsonl"))) if MOTIONS_DIR.exists() else 0, }