kassam/G1_Lootah
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
G1_Lootah/Controller/g1_mode_controller.py

698 lines
21 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env python3
"""
g1_mode_controller.py
Interactive terminal controller for Unitree G1 locomotion + basic upper body (waist/arms).
Key features:
- Menu is generated from a single OPTIONS dictionary (easy to extend).
- Ctrl+C / Quit = StopMove only (keep current posture; NO prep/damp).
- Status prints LIVE joints (deg) from LowState.
- WALK/RUN teleop uses the "old working" keyboard control (pynput).
- PREP mode = stand + balance (NO Start/FSM200).
- READY/START mode = PREP + Start (FSM 200).
- StandUp/Squat/Sit/Lie-to-Stand use SetFsmId() fallback for compatibility.
Run:
python3 g1_mode_controller.py --iface enp3s0
"""
from __future__ import annotations
import argparse
import sys
import time
from dataclasses import dataclass
from typing import Optional, Dict, Tuple, Callable, Any
from threading import Lock, Event
from collections import OrderedDict
# ---------- Unitree SDK2 ----------
from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_ # type: ignore
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
# Optional teleop keyboard:
try:
from pynput.keyboard import Listener, Key, KeyCode # type: ignore
HAVE_PYNPUT = True
except Exception:
HAVE_PYNPUT = False
# ---------------- Joint map (for status) ----------------
JOINT_NAMES = {
12: "WAIST_YAW",
13: "WAIST_ROLL",
14: "WAIST_PITCH",
15: "L_SHOULDER_PITCH",
16: "L_SHOULDER_ROLL",
17: "L_SHOULDER_YAW",
18: "L_ELBOW",
19: "L_WRIST_ROLL",
20: "L_WRIST_PITCH",
21: "L_WRIST_YAW",
22: "R_SHOULDER_PITCH",
23: "R_SHOULDER_ROLL",
24: "R_SHOULDER_YAW",
25: "R_ELBOW",
26: "R_WRIST_ROLL",
27: "R_WRIST_PITCH",
28: "R_WRIST_YAW",
}
def rad2deg(x: float) -> float:
return float(x) * 180.0 / 3.141592653589793
# ---------------- Live State Monitor ----------------
class LiveStateMonitor:
"""Subscribes LowState and keeps last joint positions."""
def __init__(self):
self.lock = Lock()
self.last_q_rad: Dict[int, float] = {}
self.msg_count = 0
self.last_rx_time = 0.0
self.first_msg_evt = Event()
def cb(self, msg: LowState_):
motor_state = getattr(msg, "motor_state", None)
if motor_state is None:
return
with self.lock:
self.msg_count += 1
self.last_rx_time = time.time()
for idx in JOINT_NAMES.keys():
try:
self.last_q_rad[idx] = float(motor_state[idx].q)
except Exception:
pass
self.first_msg_evt.set()
def snapshot_deg(self) -> Tuple[int, float, Dict[int, float]]:
with self.lock:
age = (time.time() - self.last_rx_time) if self.last_rx_time else 1e9
qdeg = {i: rad2deg(q) for i, q in self.last_q_rad.items()}
return self.msg_count, age, qdeg
# ---------------- Upper body (waist/arms) publisher ----------------
@dataclass
class ArmRig:
pub: object
cmd: object
crc: object
waist_idx: int = 12
def set_joint(self, idx: int, q: float, kp: float = 60.0, kd: float = 1.5, tau: float = 0.0):
mc = self.cmd.motor_cmd[idx]
mc.q = float(q)
mc.dq = 0.0
mc.tau = float(tau)
mc.kp = float(kp)
mc.kd = float(kd)
def send(self):
self.cmd.crc = self.crc.Crc(self.cmd)
self.pub.Write(self.cmd)
def try_init_arm_sdk() -> Optional[ArmRig]:
try:
from unitree_sdk2py.core.channel import ChannelPublisher # type: ignore
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ # type: ignore
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ # type: ignore
from unitree_sdk2py.utils.crc import CRC # type: ignore
pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
pub.Init()
cmd = unitree_hg_msg_dds__LowCmd_()
crc = CRC()
# enable arm_sdk mode
cmd.motor_cmd[29].q = 1
return ArmRig(pub=pub, cmd=cmd, crc=crc)
except Exception as e:
print("[arm_sdk] disabled:", e)
return None
# ---------------- Core helpers ----------------
def init_loco(timeout: float = 10.0) -> LocoClient:
bot = LocoClient()
bot.SetTimeout(timeout)
bot.Init()
return bot
def init_msc(timeout: float = 5.0) -> MotionSwitcherClient:
msc = MotionSwitcherClient()
msc.SetTimeout(timeout)
msc.Init()
return msc
def safe_call(name: str, fn: Callable[..., Any], *args, **kwargs) -> Any:
"""
Unitree python wrappers often return None on success.
This helper prints [OK] on no-exception; prints [ERR] on exception.
"""
try:
ret = fn(*args, **kwargs)
print(f"[OK] {name}: success (ret={ret})")
return ret
except Exception as e:
code = getattr(e, "code", None)
payload = getattr(e, "payload", None)
print(f"[ERR] {name} failed: {repr(e)} code={code} payload={payload}")
return None
def stop_only(bot: Optional[LocoClient]):
"""Stop walking/teleop velocity but do NOT change mode/posture."""
if bot is None:
return
safe_call("StopMove", bot.StopMove)
def move_cmd(bot: LocoClient, vx: float, vy: float, om: float):
"""
SDK variations exist. We try a safe call order.
"""
try:
bot.Move(vx, vy, om, True)
return
except Exception:
pass
try:
bot.Move(vx, vy, om)
return
except Exception:
pass
try:
bot.Move(vx, vy, om, continuous_move=True) # type: ignore
return
except Exception as e:
print("[ERR] Move failed in all signatures:", repr(e))
def fsm_set(bot: LocoClient, fsm_id: int, label: str):
"""
Prefer direct methods if present; fallback to SetFsmId(fsm_id).
"""
if hasattr(bot, label):
fn = getattr(bot, label)
return safe_call(label, fn)
# common fallback in some sdk builds
if hasattr(bot, "SetFsmId"):
return safe_call(f"{label} (FSM {fsm_id}) via SetFsmId({fsm_id})", bot.SetFsmId, fsm_id)
print(f"[ERR] Neither {label} nor SetFsmId exists in this SDK build.")
return None
def balance_stand(bot: LocoClient, mode: int = 0):
"""
Some builds require BalanceStand(balance_mode).
Others expose SetBalanceMode(0/1).
"""
if hasattr(bot, "BalanceStand"):
try:
return safe_call(f"BalanceStand({mode})", bot.BalanceStand, mode)
except TypeError:
return safe_call("BalanceStand()", bot.BalanceStand)
if hasattr(bot, "SetBalanceMode"):
return safe_call(f"SetBalanceMode({mode})", bot.SetBalanceMode, mode)
print("[WARN] No BalanceStand/SetBalanceMode available in this SDK.")
return None
# ---------------- PREP / READY sequences ----------------
def prep_mode(bot: LocoClient):
"""
PREP = stand + balance (NO Start/FSM200).
Matches your tested behavior:
StopMove -> Damp -> StandUp(FSM4) -> ramp stand height -> BalanceStand(0) -> set final height.
"""
print("[PREP] stand + balance (NO Start/FSM200) ...")
safe_call("StopMove", bot.StopMove)
safe_call("Damp", bot.Damp)
# StandUp FSM 4
fsm_set(bot, 4, "StandUp")
# Stand height ramp (0.02 .. 0.50)
if hasattr(bot, "SetStandHeight"):
for h in [x / 100.0 for x in range(2, 51, 2)]:
safe_call(f"SetStandHeight({h:.2f}m)", bot.SetStandHeight, float(h))
time.sleep(0.03)
# set a comfortable final height before balancing
safe_call("SetStandHeight(0.22m)", bot.SetStandHeight, 0.22)
else:
print("[WARN] bot.SetStandHeight not available in this SDK build.")
# Balance (static stand)
balance_stand(bot, mode=0)
# Re-send final height (common trick)
if hasattr(bot, "SetStandHeight"):
safe_call("SetStandHeight(0.22m)", bot.SetStandHeight, 0.22)
print("[PREP] Done. (NO Start)")
def ready_start_mode(bot: LocoClient):
"""
READY/START = PREP + Start (FSM200).
"""
print("[READY] stand + balance + start (FSM 200) ...")
prep_mode(bot)
if hasattr(bot, "Start"):
safe_call("Start (FSM 200)", bot.Start)
else:
# some builds: SetFsmId(200)
fsm_set(bot, 200, "Start")
print("[READY] Done. (FSM200 expected)")
# ---------------- Teleop (old keyboard function) ----------------
def teleop_loop(bot: LocoClient, speed_limit: float):
if not HAVE_PYNPUT:
print("[ERR] pynput not installed. Install: pip install pynput")
return
LIN_STEP = 0.05
ANG_STEP = 0.2
SEND_PERIOD = 0.10 # 10 Hz
pressed = set()
def on_press(k):
if isinstance(k, KeyCode) and k.char:
pressed.add(k.char.lower())
else:
pressed.add(k)
def on_release(k):
if isinstance(k, KeyCode) and k.char:
pressed.discard(k.char.lower())
else:
pressed.discard(k)
def key(name: str) -> bool:
if name == "space":
return Key.space in pressed
if name == "esc":
return Key.esc in pressed
return name in pressed
def clamp(v: float) -> float:
return max(-speed_limit, min(speed_limit, v))
vx = vy = om = 0.0
last_send = 0.0
print("\n--- TELEOP ---")
print("Hold keys: W/S (vx), Q/E (vy), A/D (yaw), Space (stop). ESC to exit.\n")
listener = Listener(on_press=on_press, on_release=on_release)
listener.start()
try:
while True:
if key("esc"):
stop_only(bot)
break
if key("w") and not key("s"):
vx = clamp(vx + LIN_STEP)
elif key("s") and not key("w"):
vx = clamp(vx - LIN_STEP)
else:
vx = 0.0
if key("q") and not key("e"):
vy = clamp(vy + LIN_STEP)
elif key("e") and not key("q"):
vy = clamp(vy - LIN_STEP)
else:
vy = 0.0
if key("a") and not key("d"):
om = clamp(om + ANG_STEP)
elif key("d") and not key("a"):
om = clamp(om - ANG_STEP)
else:
om = 0.0
if key("space"):
vx = vy = om = 0.0
now = time.time()
if now - last_send >= SEND_PERIOD:
# enable gait mode if supported (some SDKs require it for walking)
if hasattr(bot, "SetBalanceMode"):
try:
bot.SetBalanceMode(1)
except Exception:
pass
move_cmd(bot, vx, vy, om)
last_send = now
sys.stdout.write(f"\rvx={vx:+.2f} vy={vy:+.2f} yaw={om:+.2f} ")
sys.stdout.flush()
time.sleep(0.005)
except KeyboardInterrupt:
stop_only(bot)
print("\n--- teleop interrupted (StopMove only) ---\n")
finally:
listener.stop()
print("\n--- teleop ended ---\n")
# ---------------- Presets for arms ----------------
LEFT_DOWN = [
(12, 0.0),
(15, +0.211),
(16, +0.181),
(17, -0.284),
(18, +0.672),
(19, -0.379),
(20, -0.852),
(21, -0.019),
]
RIGHT_DOWN = [
(12, 0.0),
(22, +0.087),
(23, -0.271),
(24, +0.323),
(25, +0.691),
(26, +0.240),
(27, -0.771),
(28, -0.176),
]
def apply_pose(arm: ArmRig, pose, kp: float = 60.0, kd: float = 1.5):
for j, q in pose:
arm.set_joint(j, q, kp=kp, kd=kd, tau=0.0)
arm.send()
# ---------------- Main ----------------
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--iface", required=True, help="NIC connected to G1 (example: enp3s0)")
ap.add_argument("--timeout", type=float, default=10.0)
ap.add_argument("--topic", default="", help="LowState topic override (default tries rt/lf/lowstate then rt/lowstate)")
args = ap.parse_args()
# DDS init ONCE
ChannelFactoryInitialize(0, args.iface)
# Live monitor
monitor = LiveStateMonitor()
topics = [args.topic] if args.topic else ["rt/lf/lowstate", "rt/lowstate"]
sub = None
last_err = None
for t in topics:
try:
sub = ChannelSubscriber(t, LowState_)
sub.Init(monitor.cb, 200)
print(f"[OK] LowState subscribed: {t}")
break
except Exception as e:
last_err = e
print(f"[WARN] LowState subscribe failed for {t}: {e}", file=sys.stderr)
if sub is None:
print(f"[WARN] Could not subscribe LowState. Last error: {last_err}")
# Init clients
print("Initializing LocoClient...")
bot: Optional[LocoClient] = None
try:
bot = init_loco(timeout=args.timeout)
except Exception as e:
print("[ERR] Loco init failed:", e)
print("Initializing MotionSwitcherClient...")
msc: Optional[MotionSwitcherClient] = None
try:
msc = init_msc(timeout=5.0)
status, result = msc.CheckMode()
print(f"[MSC] Current mode: {result.get('name') or '(none)'}")
except Exception as e:
print("[WARN] MotionSwitcher init failed:", e)
arm = try_init_arm_sdk()
# Snapshot on start
if sub is not None and monitor.first_msg_evt.wait(timeout=2.0):
cnt, age, qdeg = monitor.snapshot_deg()
print("\n[STATE] Current joint snapshot (deg):")
print(f" lowstate_msgs={cnt} age={age:.2f}s")
for idx in sorted(qdeg.keys()):
print(f" [{idx:02d}] {JOINT_NAMES.get(idx,'?'):<18} {qdeg[idx]:+8.2f}°")
print()
def need_bot() -> bool:
if bot is None:
print("[ERR] No LocoClient available.")
return False
return True
def need_arm() -> bool:
if arm is None:
print("[ERR] arm_sdk not available.")
return False
return True
# ----- Dictionary-driven menu -----
OPTIONS: "OrderedDict[int, Dict[str, Any]]" = OrderedDict()
OPTIONS[1] = {"label": "PREP mode (stand + balance) [NO start/FSM200]",
"fn": lambda: prep_mode(bot) if need_bot() else None}
OPTIONS[2] = {"label": "READY/START mode (stand + balance + start, FSM 200)",
"fn": lambda: ready_start_mode(bot) if need_bot() else None}
OPTIONS[3] = {"label": "ZeroTorque (FSM 0)",
"fn": lambda: fsm_set(bot, 0, "ZeroTorque") if need_bot() else None}
OPTIONS[4] = {"label": "Damp (FSM 1)",
"fn": lambda: fsm_set(bot, 1, "Damp") if need_bot() else None}
OPTIONS[5] = {"label": "StandUp (FSM 4)",
"fn": lambda: fsm_set(bot, 4, "StandUp") if need_bot() else None}
OPTIONS[6] = {"label": "Squat (FSM 2)",
"fn": lambda: fsm_set(bot, 2, "Squat") if need_bot() else None}
OPTIONS[7] = {"label": "Sit (FSM 3)",
"fn": lambda: fsm_set(bot, 3, "Sit") if need_bot() else None}
OPTIONS[8] = {"label": "LowStand",
"fn": lambda: safe_call("LowStand", bot.LowStand) if (need_bot() and hasattr(bot, "LowStand")) else print("[ERR] LowStand not available")}
OPTIONS[9] = {"label": "HighStand",
"fn": lambda: safe_call("HighStand", bot.HighStand) if (need_bot() and hasattr(bot, "HighStand")) else print("[ERR] HighStand not available")}
OPTIONS[10] = {"label": "Lie-to-Stand (FSM 702) (if supported)",
"fn": lambda: fsm_set(bot, 702, "Lie2StandUp") if need_bot() else None}
OPTIONS[11] = {"label": "Walk teleop (keyboard: W/S Q/E A/D, ESC exit)",
"fn": lambda: teleop_loop(bot, speed_limit=0.6) if need_bot() else None}
OPTIONS[12] = {"label": "Run teleop (keyboard: W/S Q/E A/D, ESC exit)",
"fn": lambda: teleop_loop(bot, speed_limit=1.2) if need_bot() else None}
OPTIONS[16] = {"label": "Status (LIVE joints)",
"fn": lambda: show_status(bot, arm, msc, sub, monitor)}
OPTIONS[17] = {"label": "MotionSwitcher: Select mode 'ai' (supported)",
"fn": lambda: msc_select_ai(msc)}
OPTIONS[18] = {"label": "MotionSwitcher: Release current mode",
"fn": lambda: msc_release(msc)}
OPTIONS[19] = {"label": "MotionSwitcher: Show current mode",
"fn": lambda: msc_show(msc)}
OPTIONS[20] = {"label": "Seated mode ON (Sit)",
"fn": lambda: fsm_set(bot, 3, "Sit") if need_bot() else None}
OPTIONS[21] = {"label": "Seated mode OFF (StandUp + PREP)",
"fn": lambda: seated_off(bot) if need_bot() else None}
OPTIONS[22] = {"label": "Reconnect (re-init Loco + MotionSwitcher)",
"fn": lambda: reconnect(args, lambda b: set_bot_ref(b), lambda m: set_msc_ref(m))}
# Helpers need access to outer bot/msc refs:
def set_bot_ref(new_bot: Optional[LocoClient]):
nonlocal bot
bot = new_bot
def set_msc_ref(new_msc: Optional[MotionSwitcherClient]):
nonlocal msc
msc = new_msc
def print_menu():
print("\n" + "=" * 70)
print("G1 MODE CONTROLLER - select option:")
print("=" * 70)
for k, v in OPTIONS.items():
print(f"{k:>3}) {v['label']}")
print(" 0) Quit (StopMove only)")
print("=" * 70)
print_menu()
while True:
try:
choice = input("\nSelect option (0-22): ").strip()
except KeyboardInterrupt:
stop_only(bot)
print("\n[EXIT] Ctrl+C → StopMove only. Keeping current posture. Bye.")
break
except EOFError:
stop_only(bot)
print("\n[EXIT] EOF → StopMove only. Keeping current posture. Bye.")
break
if not choice:
continue
if choice == "0":
stop_only(bot)
print("\n[EXIT] Quit → StopMove only. Keeping current posture. Bye.")
break
try:
opt = int(choice)
except ValueError:
print("Please enter a number.")
continue
if opt not in OPTIONS:
print("Unknown option.")
print_menu()
continue
try:
OPTIONS[opt]["fn"]()
except KeyboardInterrupt:
stop_only(bot)
print("\n[INTERRUPT] Ctrl+C → StopMove only. Keeping current posture.")
except Exception as e:
print("[ERR] Action failed:", repr(e))
print_menu()
# --------- extra action helpers that need current refs ---------
def show_status(bot, arm, msc, sub, monitor):
print("Status:")
print(f" LocoClient: {'OK' if bot is not None else 'NOT INITIALISED'}")
print(f" MotionSwitcher: {'OK' if msc is not None else 'DISABLED'}")
if msc is not None:
try:
_s, r = msc.CheckMode()
print(f" MSC mode: {r.get('name') or '(none)'}")
except Exception:
print(" MSC mode: (error reading)")
print(f" arm_sdk: {'OK' if arm is not None else 'DISABLED'}")
print(f" pynput: {'OK' if HAVE_PYNPUT else 'NOT INSTALLED'}")
if sub is None:
print(" LowState: NOT SUBSCRIBED")
return
cnt, age, qdeg = monitor.snapshot_deg()
print(f" LowState: OK msgs={cnt} age={age:.2f}s")
print(" Current joints (deg):")
for idx in sorted(qdeg.keys()):
print(f" [{idx:02d}] {JOINT_NAMES.get(idx,'?'):<18} {qdeg[idx]:+8.2f}°")
def msc_select_ai(msc: Optional[MotionSwitcherClient]):
if msc is None:
print("[ERR] MotionSwitcher not available.")
return
ret = safe_call("SelectMode('ai')", msc.SelectMode, "ai")
# ret could be (7004,None) etc; we just print it from safe_call.
def msc_release(msc: Optional[MotionSwitcherClient]):
if msc is None:
print("[ERR] MotionSwitcher not available.")
return
safe_call("ReleaseMode()", msc.ReleaseMode)
def msc_show(msc: Optional[MotionSwitcherClient]):
if msc is None:
print("[ERR] MotionSwitcher not available.")
return
try:
s, r = msc.CheckMode()
print(f"[MSC] Current mode: {r.get('name') or '(none)'}")
except Exception as e:
print("[ERR] CheckMode failed:", repr(e))
def seated_off(bot: LocoClient):
print("[SEATED] Leaving seated: StandUp + PREP (no Start)...")
fsm_set(bot, 4, "StandUp")
prep_mode(bot)
def reconnect(args, set_bot, set_msc):
print("Reconnecting LocoClient...")
try:
b = init_loco(timeout=args.timeout)
set_bot(b)
print("[OK] LocoClient reconnected.")
except Exception as e:
print("[ERR] Loco reconnect failed:", repr(e))
set_bot(None)
print("Reconnecting MotionSwitcherClient...")
try:
m = init_msc(timeout=5.0)
set_msc(m)
try:
_s, r = m.CheckMode()
print(f"[MSC] Current mode: {r.get('name') or '(none)'}")
except Exception:
pass
print("[OK] MotionSwitcher reconnected.")
except Exception as e:
print("[ERR] MotionSwitcher reconnect failed:", repr(e))
set_msc(None)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink