kassam/Sanad
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Sanad/motion/arm_controller.py

588 lines
22 KiB
Python
Raw Blame History

"""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,
}
Reference in New Issue View Git Blame Copy Permalink