382 lines
15 KiB
Python
382 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
G1 REPLAY TRIGGER (R2 + X) - SINGLE FILE
|
||
---------------------------------------
|
||
- Hold R2 + X to replay DataG1/photo_G3.jsonl (default).
|
||
- While playing: ignore any additional R2+X until replay finishes and returns home.
|
||
- Cancel combo: R2 + L1 cancels replay immediately and returns arms to home.
|
||
|
||
Usage:
|
||
python3 g1_replay_trigger_r2x.py enp3s0
|
||
python3 g1_replay_trigger_r2x.py enp3s0 --input photo_G3.jsonl --home arm_home.jsonl --speed 1.0
|
||
"""
|
||
|
||
import time
|
||
import sys
|
||
import json
|
||
import argparse
|
||
import struct
|
||
from pathlib import Path
|
||
|
||
from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize
|
||
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
|
||
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
|
||
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
|
||
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_
|
||
from unitree_sdk2py.utils.crc import CRC
|
||
|
||
from Logger import Logs
|
||
|
||
# ---------------- Logger (exactly as requested) ----------------
|
||
controller_logs = Logs()
|
||
controller_logs.LogEngine("G1_Logs", "g1_replay_trigger_r2x.log")
|
||
|
||
# ---------------- Constants ----------------
|
||
G1_NUM_MOTOR = 29
|
||
ENABLE_ARM_SDK_INDEX = 29
|
||
|
||
DATA_DIR = Path("DataG1")
|
||
REPLAY_HZ = 60.0
|
||
|
||
# --- GAINS (same as your replay) ---
|
||
KP_HIGH = 300.0 # Core/Legs
|
||
KD_HIGH = 3.0
|
||
KP_LOW = 80.0 # Arms/Ankles
|
||
KD_LOW = 3.0
|
||
KP_WRIST = 40.0
|
||
KD_WRIST = 1.5
|
||
|
||
WEAK_MOTORS = [4, 10, 15, 16, 17, 18, 22, 23, 24, 25]
|
||
WRIST_MOTORS = [19, 20, 21, 26, 27, 28]
|
||
|
||
# ---------------- Helpers ----------------
|
||
def resolve_input_path(in_path: str) -> str:
|
||
p = Path(in_path)
|
||
if len(p.parts) == 1:
|
||
return str(DATA_DIR / p.name)
|
||
return str(p)
|
||
|
||
def load_home_pose(home_path: str):
|
||
"""Reads the last frame of arm_home.jsonl to get the target pose."""
|
||
path = resolve_input_path(home_path)
|
||
try:
|
||
last_valid_q = None
|
||
with open(path, 'r') as f:
|
||
for line in f:
|
||
d = json.loads(line)
|
||
if 'q' in d and len(d['q']) == G1_NUM_MOTOR:
|
||
last_valid_q = d['q']
|
||
if last_valid_q:
|
||
controller_logs.print_and_log(f"✅ Loaded Home Pose from {path}", "info")
|
||
return last_valid_q
|
||
else:
|
||
controller_logs.print_and_log(f"⚠️ Warning: {path} found but contained no valid 'q' data.", "warning")
|
||
except FileNotFoundError:
|
||
controller_logs.print_and_log(f"⚠️ Warning: Home file {path} not found.", "warning")
|
||
|
||
controller_logs.print_and_log("⚠️ Using Default Home (Arms at 0.0)", "warning")
|
||
return [0.0] * G1_NUM_MOTOR
|
||
|
||
# ---------------- Wireless Controller Parser (embedded) ----------------
|
||
class unitreeRemoteController:
|
||
def __init__(self):
|
||
self.Lx = 0; self.Rx = 0; self.Ry = 0; self.Ly = 0
|
||
self.L1 = 0; self.L2 = 0; self.R1 = 0; self.R2 = 0
|
||
self.A = 0; self.B = 0; self.X = 0; self.Y = 0
|
||
self.Up = 0; self.Down = 0; self.Left = 0; self.Right = 0
|
||
self.Select = 0; self.F1 = 0; self.F3 = 0; self.Start = 0
|
||
|
||
def parse_botton(self, data1, data2):
|
||
self.R1 = (data1 >> 0) & 1; self.L1 = (data1 >> 1) & 1
|
||
self.Start = (data1 >> 2) & 1; self.Select = (data1 >> 3) & 1
|
||
self.R2 = (data1 >> 4) & 1; self.L2 = (data1 >> 5) & 1
|
||
self.F1 = (data1 >> 6) & 1; self.F3 = (data1 >> 7) & 1
|
||
self.A = (data2 >> 0) & 1; self.B = (data2 >> 1) & 1
|
||
self.X = (data2 >> 2) & 1; self.Y = (data2 >> 3) & 1
|
||
self.Up = (data2 >> 4) & 1; self.Right = (data2 >> 5) & 1
|
||
self.Down = (data2 >> 6) & 1; self.Left = (data2 >> 7) & 1
|
||
|
||
def parse_key(self, data):
|
||
offsets = [4, 8, 12, 20] # Lx, Rx, Ry, Ly
|
||
self.Lx, self.Rx, self.Ry, self.Ly = [struct.unpack('<f', data[o:o+4])[0] for o in offsets]
|
||
|
||
def parse(self, remoteData):
|
||
self.parse_key(remoteData)
|
||
self.parse_botton(remoteData[2], remoteData[3])
|
||
|
||
def get_state(self):
|
||
return self.__dict__.copy()
|
||
|
||
# ---------------- Replay Engine ----------------
|
||
class ReplayWithHome:
|
||
def __init__(self, watchdog_timeout=0.25, watchdog_disable_after=1.0):
|
||
self.low_state = None
|
||
self.low_cmd = unitree_hg_msg_dds__LowCmd_()
|
||
self.crc = CRC()
|
||
|
||
self.arm_pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
|
||
self.arm_pub.Init()
|
||
|
||
self.state_sub = None
|
||
self.first_state = False
|
||
|
||
# watchdog
|
||
self.last_state_time = 0.0
|
||
self.watchdog_timeout = float(watchdog_timeout)
|
||
self.watchdog_disable_after = float(watchdog_disable_after)
|
||
|
||
# controller
|
||
self.remote = unitreeRemoteController()
|
||
self.controller_state = self.remote.get_state()
|
||
|
||
self.is_playing = False
|
||
|
||
def InitStateSubscriber(self, topic: str):
|
||
self.state_sub = ChannelSubscriber(topic, LowState_)
|
||
self.state_sub.Init(self.LowStateHandler, 10)
|
||
|
||
def LowStateHandler(self, msg: LowState_):
|
||
self.low_state = msg
|
||
self.first_state = True
|
||
self.last_state_time = time.time()
|
||
try:
|
||
self.remote.parse(msg.wireless_remote)
|
||
self.controller_state = self.remote.get_state()
|
||
except Exception:
|
||
pass
|
||
|
||
def StateFresh(self) -> bool:
|
||
return (time.time() - self.last_state_time) < self.watchdog_timeout
|
||
|
||
def CancelRequested(self) -> bool:
|
||
"""Cancel combo: R2 + L1"""
|
||
s = self.controller_state
|
||
return bool(s.get("R2", 0) and s.get("L1", 0))
|
||
|
||
def SendFrame(self, arm_target_q, body_lock_q):
|
||
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
|
||
self.low_cmd.motor_cmd[i].tau = 0
|
||
|
||
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]
|
||
|
||
if i in WEAK_MOTORS:
|
||
self.low_cmd.motor_cmd[i].kp = KP_LOW
|
||
self.low_cmd.motor_cmd[i].kd = KD_LOW
|
||
elif i in WRIST_MOTORS:
|
||
self.low_cmd.motor_cmd[i].kp = KP_WRIST
|
||
self.low_cmd.motor_cmd[i].kd = KD_WRIST
|
||
else:
|
||
self.low_cmd.motor_cmd[i].kp = KP_HIGH
|
||
self.low_cmd.motor_cmd[i].kd = KD_HIGH
|
||
|
||
self.low_cmd.crc = self.crc.Crc(self.low_cmd)
|
||
self.arm_pub.Write(self.low_cmd)
|
||
|
||
def DisableSDK(self):
|
||
controller_logs.print_and_log("🔌 Disabling SDK...", "info")
|
||
self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0
|
||
self.low_cmd.crc = self.crc.Crc(self.low_cmd)
|
||
for _ in range(10):
|
||
self.arm_pub.Write(self.low_cmd)
|
||
time.sleep(0.02)
|
||
|
||
def ReturnArmsHome(self, last_arm_q, body_lock_q, home_q, home_steps=180):
|
||
"""Always go home smoothly (arms only), body remains locked."""
|
||
controller_logs.print_and_log("🏡 Returning arms to HOME...", "info")
|
||
for k in range(home_steps):
|
||
stale_for = (time.time() - self.last_state_time)
|
||
if stale_for > self.watchdog_disable_after:
|
||
controller_logs.print_and_log("🛑 WATCHDOG: connection lost during home. Disabling SDK.", "error")
|
||
self.DisableSDK()
|
||
return
|
||
|
||
if not self.StateFresh():
|
||
controller_logs.print_and_log("⚠️ WATCHDOG: state stale during home. Holding last pose...", "warning")
|
||
self.SendFrame(last_arm_q, body_lock_q)
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
continue
|
||
|
||
alpha = k / home_steps
|
||
interp_q = list(last_arm_q)
|
||
for j in range(15, 29):
|
||
interp_q[j] = (1-alpha)*last_arm_q[j] + alpha*home_q[j]
|
||
self.SendFrame(interp_q, body_lock_q)
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
|
||
controller_logs.print_and_log("✅ Home Reached.", "info")
|
||
|
||
def RunReplay(self, filename: str, home_filename: str, speed: float):
|
||
self.is_playing = True
|
||
cancel_requested = False
|
||
|
||
try:
|
||
controller_logs.print_and_log(f"🎬 Triggered replay: {filename}", "info")
|
||
|
||
controller_logs.print_and_log("Waiting for robot...", "info")
|
||
while not self.first_state:
|
||
time.sleep(0.05)
|
||
controller_logs.print_and_log("✅ Robot Connected!", "info")
|
||
|
||
home_q = load_home_pose(home_filename)
|
||
full_body_lock_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]
|
||
|
||
frames = []
|
||
with open(filename, 'r') as f:
|
||
for line in f:
|
||
d = json.loads(line)
|
||
if 'q' in d:
|
||
frames.append(d)
|
||
|
||
if not frames:
|
||
controller_logs.print_and_log("❌ No frames found in input file.", "error")
|
||
return
|
||
|
||
controller_logs.print_and_log(f"🟢 Ready to play {len(frames)} frames.", "info")
|
||
controller_logs.print_and_log("🔒 Body LOCKED. Arms replay then return HOME.", "info")
|
||
controller_logs.print_and_log("🛑 Cancel combo: Hold R2 + L1", "info")
|
||
|
||
# Move to start
|
||
controller_logs.print_and_log("Moving to start...", "info")
|
||
file_start_q = frames[0]['q']
|
||
last_played_q = file_start_q # will update
|
||
|
||
steps = 60
|
||
for k in range(steps):
|
||
if self.CancelRequested():
|
||
controller_logs.print_and_log("🛑 CANCEL: R2+L1 detected أثناء الانتقال للبداية.", "warning")
|
||
cancel_requested = True
|
||
break
|
||
|
||
if not self.StateFresh():
|
||
controller_logs.print_and_log("⚠️ WATCHDOG: stale while moving to start. Holding...", "warning")
|
||
self.SendFrame(last_played_q, full_body_lock_q)
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
continue
|
||
|
||
alpha = k / steps
|
||
interp_q = list(full_body_lock_q)
|
||
for j in range(15, 29):
|
||
interp_q[j] = (1-alpha)*full_body_lock_q[j] + alpha*file_start_q[j]
|
||
self.SendFrame(interp_q, full_body_lock_q)
|
||
last_played_q = interp_q
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
|
||
# Play replay (unless canceled)
|
||
if not cancel_requested:
|
||
controller_logs.print_and_log("▶️ Playing...", "info")
|
||
play_elapsed = 0.0
|
||
last_real = time.time()
|
||
|
||
while True:
|
||
if self.CancelRequested():
|
||
controller_logs.print_and_log("🛑 CANCEL: R2+L1 detected أثناء التشغيل.", "warning")
|
||
cancel_requested = True
|
||
break
|
||
|
||
stale_for = (time.time() - self.last_state_time)
|
||
if stale_for > self.watchdog_disable_after:
|
||
controller_logs.print_and_log("🛑 WATCHDOG: connection lost. Disabling SDK.", "error")
|
||
self.DisableSDK()
|
||
return
|
||
|
||
if not self.StateFresh():
|
||
controller_logs.print_and_log("⚠️ WATCHDOG: timeout! Holding last pose...", "warning")
|
||
self.SendFrame(last_played_q, full_body_lock_q)
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
continue
|
||
|
||
now_real = time.time()
|
||
dt_real = now_real - last_real
|
||
last_real = now_real
|
||
play_elapsed += dt_real * speed
|
||
|
||
target_frame = None
|
||
for fr in frames:
|
||
if fr['t'] - frames[0]['t'] >= play_elapsed:
|
||
target_frame = fr
|
||
break
|
||
if target_frame is None:
|
||
break
|
||
|
||
self.SendFrame(target_frame['q'], full_body_lock_q)
|
||
last_played_q = target_frame['q']
|
||
time.sleep(1.0 / REPLAY_HZ)
|
||
|
||
# Always go home (normal finish OR cancel)
|
||
self.ReturnArmsHome(last_played_q, full_body_lock_q, home_q, home_steps=180)
|
||
self.DisableSDK()
|
||
|
||
except Exception as e:
|
||
controller_logs.print_and_log(f"❌ Exception in replay: {e}", "error")
|
||
try:
|
||
self.DisableSDK()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
self.is_playing = False
|
||
|
||
# ---------------- Trigger Loop (R2 + X) ----------------
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("iface", help="Network interface")
|
||
parser.add_argument("--input", default="photo_G3.jsonl", help="Input recording file (default: photo_G3.jsonl)")
|
||
parser.add_argument("--home", default="arm_home.jsonl", help="Home pose file")
|
||
parser.add_argument("--speed", type=float, default=1.0)
|
||
parser.add_argument("--state_topic", default="rt/lowstate", help="Lowstate topic (default: rt/lowstate)")
|
||
parser.add_argument("--watchdog", type=float, default=0.25, help="Watchdog timeout seconds")
|
||
parser.add_argument("--watchdog_disable_after", type=float, default=1.0, help="Disable SDK if stale longer than this")
|
||
args = parser.parse_args()
|
||
|
||
ChannelFactoryInitialize(0, args.iface)
|
||
|
||
engine = ReplayWithHome(
|
||
watchdog_timeout=args.watchdog,
|
||
watchdog_disable_after=args.watchdog_disable_after
|
||
)
|
||
engine.InitStateSubscriber(args.state_topic)
|
||
|
||
target_path = resolve_input_path(args.input)
|
||
|
||
controller_logs.print_and_log("✅ G1 R2+X Trigger Ready", "info")
|
||
controller_logs.print_and_log(f"🎯 Trigger file: {target_path}", "info")
|
||
controller_logs.print_and_log("🎮 Hold R2 + X to replay", "info")
|
||
controller_logs.print_and_log("🛑 Cancel while playing: Hold R2 + L1", "info")
|
||
|
||
# Rising-edge trigger + require release before next trigger
|
||
combo_prev = False
|
||
|
||
while True:
|
||
time.sleep(0.01)
|
||
|
||
if not engine.first_state:
|
||
continue
|
||
|
||
s = engine.controller_state
|
||
|
||
# Trigger combo: R2 + X
|
||
combo_now = bool(s.get("R2", 0) and s.get("X", 0))
|
||
|
||
# ✅ IGNORE while playing (and also prevent retrigger until released)
|
||
if engine.is_playing:
|
||
combo_prev = combo_now
|
||
continue
|
||
|
||
# Trigger only on rising edge (must release and press again)
|
||
if combo_now and not combo_prev:
|
||
controller_logs.print_and_log("[TRIGGER] R2 + X detected", "info")
|
||
engine.RunReplay(target_path, args.home, args.speed)
|
||
# After RunReplay returns, if user still holding, combo_prev will get updated next loop
|
||
# and won’t retrigger until they release and press again.
|
||
|
||
combo_prev = combo_now
|
||
|
||
if __name__ == "__main__":
|
||
main()
|