#!/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)