212 lines
7.4 KiB
Python
212 lines
7.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
G1 REPLAY V14 (HOME RETURN + FULL BODY LOCK)
|
|
--------------------------------------------
|
|
1. LOCK: Legs/Waist Rigid (Kp=300) to support weight.
|
|
2. REPLAY: Plays your recorded motion.
|
|
3. RETURN: Smoothly moves arms to 'arm_home.jsonl' pose at the end.
|
|
- Ignores leg data from home file (keeps them locked for balance).
|
|
- Slow 3-second transition for safety.
|
|
|
|
Usage:
|
|
python3 g1_replay_vr.py enp3s0 --input my_teleop_data.jsonl --home arm_home.jsonl
|
|
|
|
|
|
python3 g1_replay_vrpy enp3s0 --input change_battery.jsonl
|
|
"""
|
|
|
|
import time
|
|
import sys
|
|
import json
|
|
import argparse
|
|
import numpy as np
|
|
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
|
|
|
|
G1_NUM_MOTOR = 29
|
|
ENABLE_ARM_SDK_INDEX = 29
|
|
DATA_DIR = Path("DataG1")
|
|
REPLAY_HZ = 60.0
|
|
|
|
# --- GAINS (ROBOT_ARM.PY STANDARD) ---
|
|
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]
|
|
|
|
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:
|
|
print(f"✅ Loaded Home Pose from {path}")
|
|
return last_valid_q
|
|
else:
|
|
print(f"⚠️ Warning: {path} found but contained no valid 'q' data.")
|
|
except FileNotFoundError:
|
|
print(f"⚠️ Warning: Home file {path} not found.")
|
|
|
|
# Fallback: Zero Arms
|
|
print("⚠️ Using Default Home (Arms at 0.0)")
|
|
default_q = [0.0] * G1_NUM_MOTOR
|
|
return default_q
|
|
|
|
class ReplayWithHome:
|
|
def __init__(self):
|
|
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 = ChannelSubscriber("rt/lowstate", LowState_)
|
|
self.state_sub.Init(self.LowStateHandler, 10)
|
|
self.first_state = False
|
|
|
|
def LowStateHandler(self, msg: LowState_):
|
|
self.low_state = msg
|
|
self.first_state = True
|
|
|
|
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
|
|
|
|
# --- POSITIONS ---
|
|
# Arms (15-28) -> Follow Target (Replay or Home)
|
|
# Body (0-14) -> Follow Lock (Statue Mode)
|
|
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]
|
|
|
|
# --- GAINS ---
|
|
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 # 300.0
|
|
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):
|
|
print("\n🔌 Disabling SDK...")
|
|
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 Run(self, filename: str, home_filename: str, speed: float):
|
|
print("Waiting for robot...", end="", flush=True)
|
|
while not self.first_state: time.sleep(0.1)
|
|
print(" Connected!")
|
|
|
|
# 1. LOAD DATA
|
|
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 = []
|
|
try:
|
|
with open(filename, 'r') as f:
|
|
for line in f:
|
|
d = json.loads(line)
|
|
if 'q' in d: frames.append(d)
|
|
except Exception as e: print(f"Error: {e}"); return
|
|
|
|
print(f"🟢 Ready to play {len(frames)} frames.")
|
|
print(f"🔒 Body is LOCKED. Arms will return to 'arm_home' at end.")
|
|
input("👉 Press Enter to Begin...")
|
|
|
|
# 2. MOVE TO START
|
|
print("Moving to start...")
|
|
file_start_q = frames[0]['q']
|
|
steps = 60
|
|
for k in range(steps):
|
|
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)
|
|
time.sleep(1.0/REPLAY_HZ)
|
|
|
|
# 3. PLAY REPLAY
|
|
print("▶️ Playing...")
|
|
last_played_q = file_start_q
|
|
try:
|
|
t_start = time.time()
|
|
while True:
|
|
now = (time.time() - t_start) * speed
|
|
target_frame = None
|
|
for f in frames:
|
|
if f['t'] - frames[0]['t'] >= now:
|
|
target_frame = f
|
|
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)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nStopped.")
|
|
|
|
# 4. RETURN TO ARM HOME (Slow & Smooth)
|
|
print(f"\n🏡 Returning arms to {home_filename}...")
|
|
|
|
# 3 Seconds duration for smoothness
|
|
home_steps = 180
|
|
|
|
for k in range(home_steps):
|
|
alpha = k / home_steps
|
|
# Interpolate: Last Pose -> Home Pose
|
|
interp_q = list(last_played_q)
|
|
for j in range(15, 29):
|
|
interp_q[j] = (1-alpha)*last_played_q[j] + alpha*home_q[j]
|
|
|
|
# Send (Body still locked to original standing pose)
|
|
self.SendFrame(interp_q, full_body_lock_q)
|
|
time.sleep(1.0/REPLAY_HZ)
|
|
|
|
print("✅ Home Reached.")
|
|
self.DisableSDK()
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("iface", help="Network interface")
|
|
parser.add_argument("--input", required=True, help="Input recording file")
|
|
parser.add_argument("--home", default="arm_home.jsonl", help="Home pose file")
|
|
parser.add_argument("--speed", type=float, default=1.0)
|
|
args = parser.parse_args()
|
|
|
|
ChannelFactoryInitialize(0, args.iface)
|
|
path = resolve_input_path(args.input)
|
|
ReplayWithHome().Run(path, args.home, args.speed) |